@marimo-team/frontend 0.25.3-dev5 → 0.25.3-dev55

package/src/core/websocket/__tests__/{useMarimoKernelConnection.test.ts → connection-state.test.ts}

@@ -3,7 +3,7 @@
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { Logger } from "@/utils/Logger";
 import { WebSocketClosedReason, ConnectionState } from "../types";
-import { classifyCloseEvent } from "../useMarimoKernelConnection";
+import { classifyCloseEvent } from "../connection-state";
 
 function classify(reason: string | undefined) {
   return classifyCloseEvent({ reason });
@@ -13,8 +13,10 @@ describe("classifyCloseEvent", () => {
   describe("transient closes (default branch)", () => {
     it("retries on empty/undefined reason", () => {
       const decision = classify(undefined);
-      expect(decision.kind).toBe("retry");
-      expect(decision.status).toEqual({ state: ConnectionState.CONNECTING });
+      expect(decision).toEqual({
+        kind: "retry",
+        status: { state: ConnectionState.CONNECTING },
+      });
     });
 
     it("treats unknown reason strings as transient and logs a warning", () => {
@@ -40,29 +42,25 @@ describe("classifyCloseEvent", () => {
       // terminal or a rejected client would reconnect forever.
       "MARIMO_UNAUTHORIZED",
       "MARIMO_KIOSK_NOT_ALLOWED",
-    ])("%s → terminal with KERNEL_DISCONNECTED, closes transport", (reason) => {
+    ])("%s → terminal with KERNEL_DISCONNECTED", (reason) => {
       const decision = classify(reason);
-      expect(decision.kind).toBe("terminal");
-      expect(decision.status).toMatchObject({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-      });
       expect(decision).toMatchObject({
         kind: "terminal",
-        closeTransport: true,
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+        },
       });
     });
 
-    it("MARIMO_KERNEL_STARTUP_ERROR → terminal + closeTransport", () => {
+    it("MARIMO_KERNEL_STARTUP_ERROR → terminal", () => {
       const decision = classify("MARIMO_KERNEL_STARTUP_ERROR");
-      expect(decision.kind).toBe("terminal");
-      expect(decision.status).toMatchObject({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
-      });
       expect(decision).toMatchObject({
         kind: "terminal",
-        closeTransport: true,
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
+        },
       });
     });
   });

package/src/core/websocket/__tests__/{useWebSocket.test.ts → createConnectionTransport.test.ts}

@@ -4,7 +4,7 @@ import { afterEach, describe, expect, it, vi } from "vitest";
 import { BasicTransport } from "../transports/basic";
 import { SseTransport } from "../transports/sse";
 import { WsTransport } from "../transports/ws";
-import { createConnectionTransport } from "../useWebSocket";
+import { createConnectionTransport } from "../transports/create";
 
 vi.mock("../../wasm/utils", () => ({
   isWasm: () => false,

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

@@ -0,0 +1,149 @@
+/* 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";
+import { fixtureMessages } from "../../host/__tests__/fixtures";
+import { agentAttachmentsAtom } from "../../attachments/state";
+import { notebookClientAtom } from "../state";
+
+vi.mock("../transports/create", () => ({ createConnectionTransport: vi.fn() }));
+
+function makeTransport(url?: () => string) {
+  const events = new EventTarget();
+  const transport: IConnectionTransport = {
+    readyState: WebSocket.CLOSED,
+    reconnect: vi.fn(() => {
+      url?.();
+      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,
+    receive: (message: unknown) =>
+      events.dispatchEvent(
+        new MessageEvent("message", { data: JSON.stringify(message) }),
+      ),
+    notify: () =>
+      events.dispatchEvent(
+        new MessageEvent("message", {
+          data: JSON.stringify({
+            event: "marimo",
+            data: { marimo: { op: "reconnected" } },
+          }),
+        }),
+      ),
+  };
+}
+
+function setup(strict = false) {
+  const store = createStore();
+  const options = {
+    static: false,
+    transportType: "websocket" as const,
+    url: () => "ws://localhost/ws",
+    headers: () => ({}),
+    waitUntilReady: async () => undefined,
+    checkHealth: async () => true,
+    onNotification: 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("exposes host presence to the UI without waiting for kernel-ready", async () => {
+    const transport = makeTransport();
+    vi.mocked(createConnectionTransport).mockReturnValue(transport.transport);
+    const { result, store, unmount } = setup();
+    await act(async () => Promise.resolve());
+    expect(store.get(notebookClientAtom)).toBe(result.current);
+    act(() =>
+      fixtureMessages("notebook-events.sse").forEach(transport.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(notebookClientAtom)).toBeNull();
+    expect(store.get(agentAttachmentsAtom)).toEqual([]);
+  });
+
+  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("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();
+  });
+});