@marimo-team/frontend 0.25.3-dev2 → 0.25.3-dev20

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

@@ -0,0 +1,140 @@
+/* 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 { ConnectionState } from "../types";
+import { createConnectionTransport } from "../transports/create";
+import type { IConnectionTransport } from "../transports/transport";
+import { useNotebookClient } from "../useNotebookClient";
+
+vi.mock("../transports/create", () => ({ createConnectionTransport: vi.fn() }));
+
+function makeTransport() {
+  const events = new EventTarget();
+  const transport: IConnectionTransport = {
+    readyState: WebSocket.CLOSED,
+    reconnect: vi.fn(() => {
+      transport.readyState = WebSocket.OPEN;
+      events.dispatchEvent(new Event("open"));
+    }),
+    close: vi.fn(() => {
+      transport.readyState = WebSocket.CLOSED;
+      events.dispatchEvent(new CloseEvent("close"));
+    }),
+    send: vi.fn(),
+    addEventListener: events.addEventListener.bind(events),
+    removeEventListener: events.removeEventListener.bind(events),
+  };
+  return {
+    transport,
+    notify: () =>
+      events.dispatchEvent(
+        new MessageEvent("message", {
+          data: JSON.stringify({
+            event: "marimo",
+            data: { marimo: { op: "reconnected" } },
+          }),
+        }),
+      ),
+  };
+}
+
+function setup(strict = false) {
+  const store = createStore();
+  const options = {
+    sessionId: "first",
+    static: false,
+    transportType: "websocket" as const,
+    messageFormat: "host" as const,
+    url: () => "ws://localhost/ws",
+    headers: () => ({}),
+    waitUntilReady: async () => undefined,
+    checkHealth: async () => true,
+    onNotification: vi.fn(),
+    onMessageError: vi.fn(),
+  };
+  const wrapper = ({ children }: PropsWithChildren) => (
+    <Provider store={store}>{children}</Provider>
+  );
+  const hook = renderHook((props) => useNotebookClient(props), {
+    initialProps: options,
+    reactStrictMode: strict,
+    wrapper,
+  });
+  return { ...hook, options, store };
+}
+
+beforeEach(() => {
+  vi.mocked(createConnectionTransport).mockReset();
+});
+
+describe("useNotebookClient", () => {
+  it("uses current callbacks and URL providers without reconnecting on a rerender", async () => {
+    const first = makeTransport();
+    vi.mocked(createConnectionTransport).mockReturnValue(first.transport);
+    const { rerender, options, unmount } = setup();
+    await act(async () => Promise.resolve());
+    const onNotification = vi.fn();
+    rerender({ ...options, onNotification, url: () => "ws://redirected/ws" });
+    act(() => first.notify());
+    expect(options.onNotification).not.toHaveBeenCalled();
+    expect(onNotification).toHaveBeenCalledExactlyOnceWith({
+      data: { op: "reconnected" },
+    });
+    expect(createConnectionTransport).toHaveBeenCalledOnce();
+    expect(vi.mocked(createConnectionTransport).mock.calls[0][0].url()).toBe(
+      "ws://redirected/ws",
+    );
+    unmount();
+  });
+
+  it("replaces the transport when session identity changes and ignores old messages", async () => {
+    const first = makeTransport();
+    const second = makeTransport();
+    vi.mocked(createConnectionTransport)
+      .mockReturnValueOnce(first.transport)
+      .mockReturnValueOnce(second.transport);
+    const { rerender, options, unmount } = setup();
+    await act(async () => Promise.resolve());
+    rerender({ ...options, sessionId: "second" });
+    await act(async () => Promise.resolve());
+    expect(first.transport.close).toHaveBeenCalledOnce();
+    act(() => {
+      first.notify();
+      second.notify();
+    });
+    expect(options.onNotification).toHaveBeenCalledOnce();
+    unmount();
+  });
+
+  it("cleans up every transport under Strict Mode and only connects the active lifetime", async () => {
+    const transports: ReturnType<typeof makeTransport>[] = [];
+    vi.mocked(createConnectionTransport).mockImplementation(() => {
+      const next = makeTransport();
+      transports.push(next);
+      return next.transport;
+    });
+    const { result, store, options, unmount } = setup(true);
+    await act(async () => Promise.resolve());
+    expect(transports).toHaveLength(2);
+    expect(transports[0].transport.reconnect).not.toHaveBeenCalled();
+    expect(transports[0].transport.close).toHaveBeenCalledOnce();
+    expect(transports[1].transport.reconnect).toHaveBeenCalledOnce();
+    expect(store.get(result.current.connectionAtom)).toEqual({
+      state: ConnectionState.OPEN,
+    });
+    act(() => {
+      for (const transport of transports) {
+        transport.notify();
+      }
+    });
+    expect(options.onNotification).toHaveBeenCalledOnce();
+    unmount();
+    expect(transports[1].transport.close).toHaveBeenCalledOnce();
+    await result.current.reconnect();
+    expect(transports[1].transport.reconnect).toHaveBeenCalledOnce();
+  });
+});

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

@@ -21,14 +21,9 @@ import {
   vi,
 } from "vitest";
 
-vi.mock("@/core/websocket/useWebSocket", async () => {
-  const actual =
-    await vi.importActual<typeof import("../useWebSocket")>("../useWebSocket");
-  return {
-    ...actual,
-    useConnectionTransport: vi.fn(),
-  };
-});
+vi.mock("../transports/create", () => ({
+  createConnectionTransport: vi.fn(),
+}));
 
 vi.mock("@/core/runtime/config", async () => {
   const actual = await vi.importActual<typeof import("@/core/runtime/config")>(
@@ -41,21 +36,33 @@ vi.mock("@/core/runtime/config", async () => {
 });
 
 import { MockNotebook } from "@/__mocks__/notebook";
+import * as Toast from "@/components/ui/use-toast";
+import { Logger } from "@/utils/Logger";
+import { Deferred } from "@/utils/Deferred";
 import { cellId } from "@/__tests__/branded";
 import { alertAtom, getPackageAlert } from "@/core/alerts/state";
+import {
+  type Attachment,
+  attachmentsAtom,
+  agentAttachmentsAtom,
+} from "@/core/attachments/state";
 import { notebookAtom } from "@/core/cells/cells";
 import { AppConfigSchema } from "@/core/config/config-schema";
 import { ConnectionNotice } from "@/components/editor/alerts/connection-notice";
 import { kernelStartupErrorAtom } from "@/core/errors/state";
+import type { RunEvent } from "@/core/host/protocol";
 import type { NotificationPayload } from "@/core/kernel/messages";
 import { useRuntimeManager } from "@/core/runtime/config";
 import { initialRunCompletedAtom } from "../../kernel/state";
 import { connectionAtom, startupProgressAtom } from "../../network/connection";
-import type { SessionId } from "../../kernel/session";
+import { generateSessionId } from "../../kernel/session";
+import { RuntimeManager } from "../../runtime/runtime";
 import { WebSocketClosedReason, ConnectionState } from "../types";
 import type { IConnectionTransport } from "../transports/transport";
-import { useMarimoKernelConnection } from "../useMarimoKernelConnection";
-import { useConnectionTransport } from "../useWebSocket";
+import { useNotebookSession } from "../useNotebookSession";
+import { createConnectionTransport } from "../transports/create";
+
+const SESSION_ID = generateSessionId();
 
 interface MockTransport {
   readyState: 0 | 1 | 2 | 3;
@@ -79,16 +86,47 @@ function makeTransport(
   };
 }
 
+function getTransportHandlers() {
+  const result = vi.mocked(createConnectionTransport).mock.results.at(-1);
+  if (!result || result.type !== "return") {
+    throw new Error("No connection transport was created");
+  }
+  const transport = result.value;
+  function emit(
+    event: "open" | "close" | "error" | "message",
+    message: Event | CloseEvent | MessageEvent,
+  ) {
+    const callback = vi
+      .mocked(transport.addEventListener)
+      .mock.calls.find(([name]) => name === event)?.[1];
+    callback?.(message);
+  }
+  return {
+    onError: (event: Event) => emit("error", event),
+    onOpen: (event: Event) => emit("open", event),
+    onClose: (event: CloseEvent) => emit("close", event),
+    onMessage: (event: MessageEvent) => emit("message", event),
+  };
+}
+
 function makeRuntimeManager(
   reconcileFromHealth = vi.fn<() => Promise<boolean>>().mockResolvedValue(true),
 ) {
-  return {
+  const runtime = new RuntimeManager(
+    { url: "http://localhost", lazy: true },
+    true,
+  );
+  vi.spyOn(runtime, "reconcileFromHealth").mockImplementation(
     reconcileFromHealth,
-    probeHealth: vi.fn().mockResolvedValue(true),
-    getWsURL: () => new URL("ws://localhost/ws"),
-    waitForHealthy: vi.fn().mockResolvedValue(undefined),
-    isSameOrigin: true,
-  };
+  );
+  vi.spyOn(runtime, "probeHealth").mockResolvedValue(true);
+  vi.spyOn(runtime, "getWsURL").mockReturnValue(new URL("ws://localhost/ws"));
+  // Keep the initial connection pending; these tests exercise session retries.
+  vi.spyOn(runtime, "waitForHealthy").mockReturnValue(
+    new Deferred<void>().promise,
+  );
+  vi.spyOn(runtime, "isSameOrigin", "get").mockReturnValue(false);
+  return runtime;
 }
 
 function renderConnectionHook(store: ReturnType<typeof createStore>) {
@@ -99,16 +137,16 @@ function renderConnectionHook(store: ReturnType<typeof createStore>) {
   );
   return renderHook(
     () =>
-      useMarimoKernelConnection({
-        sessionId: "test-session" as SessionId,
+      useNotebookSession({
+        sessionId: SESSION_ID,
         autoInstantiate: false,
-        setCells: () => {},
+        setCells: vi.fn(),
       }),
     { wrapper },
   );
 }
 
-describe("useMarimoKernelConnection.reconnect()", () => {
+describe("useNotebookSession.reconnect()", () => {
   let transport: MockTransport;
   let isHealthy: Mock<() => Promise<boolean>>;
   let store: ReturnType<typeof createStore>;
@@ -122,12 +160,8 @@ describe("useMarimoKernelConnection.reconnect()", () => {
       code: WebSocketClosedReason.KERNEL_DISCONNECTED,
       reason: "kernel not found",
     });
-    vi.mocked(useConnectionTransport).mockReturnValue(transport);
-    vi.mocked(useRuntimeManager).mockReturnValue(
-      makeRuntimeManager(isHealthy) as unknown as ReturnType<
-        typeof useRuntimeManager
-      >,
-    );
+    vi.mocked(createConnectionTransport).mockReturnValue(transport);
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager(isHealthy));
   });
 
   function renderUseHook() {
@@ -183,31 +217,146 @@ describe("useMarimoKernelConnection.reconnect()", () => {
   });
 });
 
-describe("useMarimoKernelConnection messages", () => {
-  it("records completion of the initial run", () => {
-    const store = createStore();
-    vi.mocked(useConnectionTransport).mockClear();
-    vi.mocked(useConnectionTransport).mockReturnValue(
+describe("useNotebookSession messages", () => {
+  let store: ReturnType<typeof createStore>;
+
+  beforeEach(() => {
+    store = createStore();
+    vi.mocked(createConnectionTransport).mockClear();
+    vi.mocked(createConnectionTransport).mockReturnValue(
       makeTransport(WebSocket.OPEN),
     );
-    vi.mocked(useRuntimeManager).mockReturnValue(
-      makeRuntimeManager() as unknown as ReturnType<typeof useRuntimeManager>,
-    );
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
     renderConnectionHook(store);
+  });
 
-    const options = vi.mocked(useConnectionTransport).mock.calls.at(-1)?.[0];
-    expect(store.get(initialRunCompletedAtom)).toBe(false);
+  function receive(payload: unknown) {
+    const options = getTransportHandlers();
     act(() => {
-      options?.onMessage(
-        new MessageEvent("message", {
-          data: JSON.stringify({
-            op: "completed-run",
-            data: { op: "completed-run", run_id: null },
-          }),
-        }),
+      options.onMessage(
+        new MessageEvent("message", { data: JSON.stringify(payload) }),
       );
     });
-    expect(store.get(initialRunCompletedAtom)).toBe(true);
+  }
+
+  it("replaces attachment snapshots, including the last detach", () => {
+    const agent: Attachment = {
+      id: "a1",
+      kind: "agent",
+      name: "Pi",
+      since: "2026-10-14T00:00:00Z",
+      notebook_id: "nb-test",
+      identity: "pi-thread",
+    };
+    const client: Attachment = {
+      id: "c1",
+      kind: "client",
+      name: null,
+      since: "2026-10-14T00:00:01Z",
+      notebook_id: "nb-test",
+    };
+    for (const attachments of [[agent, client], [client], []]) {
+      receive({ data: { op: "attachments", attachments } });
+      expect(store.get(attachmentsAtom)).toEqual(attachments);
+    }
+  });
+
+  it.each(["{broken", "null"])(
+    "reports a malformed legacy frame and continues applying notifications: %s",
+    (data) => {
+      const toast = vi.spyOn(Toast, "toast");
+      try {
+        receive({ data: { op: "reconnected" } });
+        act(() =>
+          getTransportHandlers().onMessage(
+            new MessageEvent("message", { data }),
+          ),
+        );
+        expect(toast).toHaveBeenCalledExactlyOnceWith({
+          title: "Failed to handle message",
+          description: expect.any(String),
+          variant: "danger",
+        });
+        expect(store.get(connectionAtom)).toEqual({
+          state: ConnectionState.OPEN,
+        });
+        receive({
+          data: {
+            op: "completed-run",
+            run_id: "r1",
+            outcome: "succeeded",
+            error: null,
+          },
+        });
+        expect(store.get(initialRunCompletedAtom)).toBe(true);
+      } finally {
+        toast.mockRestore();
+      }
+    },
+  );
+
+  it.each(["session", "host"])(
+    "records completion of the initial run from the %s stream",
+    (stream) => {
+      const notification = {
+        op: "completed-run",
+        run_id: null,
+        outcome: "succeeded",
+        error: null,
+      } satisfies NotificationPayload["data"];
+      expect(store.get(initialRunCompletedAtom)).toBe(false);
+      receive(
+        stream === "host"
+          ? { id: "42", event: "marimo", data: { marimo: notification } }
+          : { op: notification.op, data: notification },
+      );
+      expect(store.get(initialRunCompletedAtom)).toBe(true);
+    },
+  );
+
+  it("ignores other host events before handling the next notification", () => {
+    const errors = vi.spyOn(Logger, "error").mockImplementation(vi.fn());
+    try {
+      receive({ id: "1", event: "ready", data: {} });
+      receive({
+        id: "2",
+        event: "run",
+        data: {
+          run: {
+            id: "r1",
+            notebook_id: "n1",
+            caused_by: null,
+            status: "completed",
+            outcome: "succeeded",
+            started_at: "2026-10-19T12:00:00Z",
+            completed_at: "2026-10-19T12:00:01Z",
+          },
+        },
+      } satisfies RunEvent);
+      receive({
+        id: "3",
+        event: "future-event",
+        data: { marimo: { op: "completed-run", run_id: "r1" } },
+      });
+      receive({ id: "4", event: "another-event", data: null });
+      expect(store.get(initialRunCompletedAtom)).toBe(false);
+      receive({
+        id: "5",
+        event: "marimo",
+        data: {
+          marimo: {
+            op: "completed-run",
+            run_id: "r1",
+            outcome: "succeeded",
+            error: null,
+          },
+        },
+      });
+      expect(store.get(initialRunCompletedAtom)).toBe(true);
+      expect(errors).not.toHaveBeenCalled();
+    } finally {
+      errors.mockRestore();
+    }
   });
 });
 
@@ -216,15 +365,13 @@ it.each(["kernel-ready", "reconnected"])(
   async (op) => {
     const store = createStore();
     store.set(connectionAtom, { state: ConnectionState.CONNECTING });
-    vi.mocked(useConnectionTransport).mockClear();
-    vi.mocked(useConnectionTransport).mockReturnValue(
+    vi.mocked(createConnectionTransport).mockClear();
+    vi.mocked(createConnectionTransport).mockReturnValue(
       makeTransport(WebSocket.OPEN),
     );
-    vi.mocked(useRuntimeManager).mockReturnValue(
-      makeRuntimeManager() as unknown as ReturnType<typeof useRuntimeManager>,
-    );
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
     renderConnectionHook(store);
-    const options = vi.mocked(useConnectionTransport).mock.calls.at(-1)![0];
+    const options = getTransportHandlers();
 
     await act(async () => {
       await options.onOpen(new Event("open"));
@@ -261,6 +408,150 @@ it.each(["kernel-ready", "reconnected"])(
   },
 );
 
+describe("attachment snapshots across browser connections", () => {
+  const agent: Attachment = {
+    id: "a1",
+    kind: "agent",
+    name: "Pi",
+    since: "2026-10-14T00:00:00Z",
+    notebook_id: "nb-test",
+    identity: "pi-thread",
+  };
+
+  function connectedSession() {
+    const store = createStore();
+    store.set(connectionAtom, { state: ConnectionState.OPEN });
+    store.set(attachmentsAtom, [agent]);
+    const transport = makeTransport(WebSocket.OPEN);
+    vi.mocked(createConnectionTransport).mockReturnValue(transport);
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
+    renderConnectionHook(store);
+    const options = getTransportHandlers();
+    const receive = (data: NotificationPayload["data"]) =>
+      act(() => {
+        options.onMessage(
+          new MessageEvent("message", {
+            data: JSON.stringify({ op: data.op, data }),
+          }),
+        );
+      });
+    return { store, transport, options, receive };
+  }
+
+  it.each([
+    "MARIMO_NO_FILE_KEY",
+    "MARIMO_NO_SESSION_ID",
+    "MARIMO_NO_SESSION",
+    "MARIMO_SHUTDOWN",
+    "MARIMO_KERNEL_STARTUP_ERROR",
+    "MARIMO_UNAUTHORIZED",
+    "MARIMO_KIOSK_NOT_ALLOWED",
+    "MARIMO_TRANSPORT_EXHAUSTED",
+  ])("clears cached attachments on %s", (reason) => {
+    const { store, transport, options } = connectedSession();
+    act(() => options.onClose(new CloseEvent("close", { reason })));
+    expect(store.get(attachmentsAtom)).toEqual([]);
+    expect(store.get(agentAttachmentsAtom)).toEqual([]);
+    expect(store.get(connectionAtom).state).toBe(ConnectionState.CLOSED);
+    expect(transport.reconnect).not.toHaveBeenCalled();
+  });
+
+  it("clears the snapshot on transport error", () => {
+    const { store, options } = connectedSession();
+    act(() => options.onError(new Event("error")));
+    expect(store.get(attachmentsAtom)).toEqual([]);
+    expect(store.get(agentAttachmentsAtom)).toEqual([]);
+  });
+
+  it("waits for a fresh snapshot after a transient reconnect", async () => {
+    const { store, transport, options, receive } = connectedSession();
+    act(() => options.onClose(new CloseEvent("close")));
+    expect(store.get(attachmentsAtom)).toEqual([]);
+    expect(store.get(connectionAtom).state).toBe(ConnectionState.CONNECTING);
+    expect(transport.reconnect).toHaveBeenCalledOnce();
+
+    await act(async () => options.onOpen(new Event("open")));
+    receive({ op: "reconnected" });
+    expect(store.get(connectionAtom).state).toBe(ConnectionState.OPEN);
+    expect(store.get(agentAttachmentsAtom)).toEqual([]);
+
+    const refreshed = { ...agent, id: "a2", name: "Codex" };
+    receive({ op: "attachments", attachments: [refreshed] });
+    expect(store.get(agentAttachmentsAtom)).toEqual([refreshed]);
+    receive({ op: "attachments", attachments: [] });
+    expect(store.get(agentAttachmentsAtom)).toEqual([]);
+  });
+
+  it("keeps a fresh snapshot hidden until the session is ready", async () => {
+    const { store, options, receive } = connectedSession();
+    act(() => options.onClose(new CloseEvent("close")));
+    await act(async () => options.onOpen(new Event("open")));
+    receive({ op: "attachments", attachments: [agent] });
+    expect(store.get(attachmentsAtom)).toEqual([agent]);
+    expect(store.get(agentAttachmentsAtom)).toEqual([]);
+    receive({ op: "reconnected" });
+    expect(store.get(agentAttachmentsAtom)).toEqual([agent]);
+  });
+
+  it("invalidates an old snapshot when a new transport opens", async () => {
+    const { store, options } = connectedSession();
+    await act(async () => options.onOpen(new Event("open")));
+    expect(store.get(attachmentsAtom)).toEqual([]);
+    expect(store.get(agentAttachmentsAtom)).toEqual([]);
+  });
+});
+
+it.each([true, false])(
+  "keeps attachments only when kernel-ready resumes the session: %s",
+  (resumed) => {
+    const store = createStore();
+    const attachments: Attachment[] = [
+      {
+        id: "a1",
+        kind: "agent",
+        name: "Pi",
+        since: "2026-10-14T00:00:00Z",
+        notebook_id: "nb-test",
+        identity: "pi-thread",
+      },
+    ];
+    store.set(attachmentsAtom, attachments);
+    vi.mocked(createConnectionTransport).mockReturnValue(
+      makeTransport(WebSocket.OPEN),
+    );
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
+    renderConnectionHook(store);
+    const options = getTransportHandlers();
+    act(() =>
+      options.onMessage(
+        new MessageEvent("message", {
+          data: JSON.stringify({
+            op: "kernel-ready",
+            data: {
+              op: "kernel-ready",
+              cell_ids: [],
+              codes: [],
+              names: [],
+              configs: [],
+              layout: null,
+              resumed,
+              ui_values: {},
+              last_executed_code: {},
+              last_execution_time: {},
+              app_config: { width: "normal" },
+              kiosk: false,
+              capabilities: { terminal: false },
+              auto_instantiated: true,
+              consumer_capabilities: { edit: true, interact: true },
+            },
+          }),
+        }),
+      ),
+    );
+    expect(store.get(attachmentsAtom)).toEqual(resumed ? attachments : []);
+  },
+);
+
 describe("connection notice", () => {
   beforeEach(() => vi.useFakeTimers());
   afterEach(() => vi.useRealTimers());
@@ -269,15 +560,13 @@ describe("connection notice", () => {
     const store = createStore();
     store.set(connectionAtom, { state: ConnectionState.CONNECTING });
     const transport = makeTransport(WebSocket.OPEN);
-    vi.mocked(useConnectionTransport).mockReturnValue(transport);
-    vi.mocked(useRuntimeManager).mockReturnValue(
-      makeRuntimeManager() as unknown as ReturnType<typeof useRuntimeManager>,
-    );
+    vi.mocked(createConnectionTransport).mockReturnValue(transport);
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
     const Connection = () => {
-      const { reconnect } = useMarimoKernelConnection({
-        sessionId: "test-session" as SessionId,
+      const { reconnect } = useNotebookSession({
+        sessionId: SESSION_ID,
         autoInstantiate: false,
-        setCells: () => {},
+        setCells: vi.fn(),
       });
       return (
         <ConnectionNotice
@@ -293,7 +582,7 @@ describe("connection notice", () => {
         </ErrorBoundary>
       </JotaiProvider>,
     );
-    const options = vi.mocked(useConnectionTransport).mock.calls.at(-1)![0];
+    const options = getTransportHandlers();
     const send = (data: NotificationPayload["data"]) =>
       act(() => {
         options.onMessage(
@@ -482,15 +771,13 @@ describe("connection notice", () => {
 
 it("replaces environment state from snapshots, then continues live progress", () => {
   const store = createStore();
-  vi.mocked(useConnectionTransport).mockClear();
-  vi.mocked(useConnectionTransport).mockReturnValue(
+  vi.mocked(createConnectionTransport).mockClear();
+  vi.mocked(createConnectionTransport).mockReturnValue(
     makeTransport(WebSocket.OPEN),
   );
-  vi.mocked(useRuntimeManager).mockReturnValue(
-    makeRuntimeManager() as unknown as ReturnType<typeof useRuntimeManager>,
-  );
+  vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
   renderConnectionHook(store);
-  const options = vi.mocked(useConnectionTransport).mock.calls.at(-1)?.[0];
+  const options = getTransportHandlers();
   function receive(data: NotificationPayload["data"]) {
     act(() =>
       options?.onMessage(