@marimo-team/frontend 0.25.2 → 0.25.3-dev100

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/notebook", () => ({
+  createNotebookTransport: vi.fn(),
+}));
 
 vi.mock("@/core/runtime/config", async () => {
   const actual = await vi.importActual<typeof import("@/core/runtime/config")>(
@@ -41,6 +36,7 @@ vi.mock("@/core/runtime/config", async () => {
 });
 
 import { MockNotebook } from "@/__mocks__/notebook";
+import { Deferred } from "@/utils/Deferred";
 import { cellId } from "@/__tests__/branded";
 import { alertAtom, getPackageAlert } from "@/core/alerts/state";
 import { notebookAtom } from "@/core/cells/cells";
@@ -51,44 +47,96 @@ 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";
-
-interface MockTransport {
-  readyState: 0 | 1 | 2 | 3;
-  reconnect: Mock<IConnectionTransport["reconnect"]>;
-  close: Mock<IConnectionTransport["close"]>;
-  send: Mock<IConnectionTransport["send"]>;
-  addEventListener: Mock<IConnectionTransport["addEventListener"]>;
-  removeEventListener: Mock<IConnectionTransport["removeEventListener"]>;
+import { classifyCloseEvent } from "../connection-state";
+import { createNotebookTransport } from "../transports/notebook";
+import { useNotebookSession } from "../useNotebookSession";
+import { type NotebookTransport, StreamStopped } from "@marimo-team/client";
+
+const SESSION_ID = generateSessionId();
+
+/** One connection the fake transport opened. */
+interface Attempt {
+  open(): void;
+  receive(message: unknown): void;
+  drop(): void;
+  stop(error: StreamStopped): void;
 }
 
-function makeTransport(
-  readyState: 0 | 1 | 2 | 3 = WebSocket.CLOSED,
-): MockTransport {
+/** A transport the test drives, one connection at a time. */
+function makeTransport() {
+  const attempts: Attempt[] = [];
+  const transport: NotebookTransport = ({ signal }, handlers) =>
+    new Promise((resolve, reject) => {
+      attempts.push({
+        open: handlers.open,
+        receive: handlers.message,
+        drop: resolve,
+        stop: reject,
+      });
+      signal.addEventListener("abort", () => resolve(), { once: true });
+    });
+  return { transport, attempts };
+}
+
+type FakeTransport = ReturnType<typeof makeTransport>;
+
+/** Use a fake transport for the next session the hook follows. */
+function withTransport(fake: FakeTransport = makeTransport()) {
+  vi.mocked(createNotebookTransport).mockReturnValue(fake.transport);
+  return fake;
+}
+
+/** Let the stream open its first connection. */
+const started = () => act(() => Promise.resolve());
+
+/** Drive the latest connection the fake transport opened. */
+function connection(fake: FakeTransport) {
+  const current = () => {
+    const attempt = fake.attempts.at(-1);
+    if (attempt === undefined) {
+      throw new Error("No connection was opened");
+    }
+    return attempt;
+  };
   return {
-    readyState,
-    reconnect: vi.fn(),
-    close: vi.fn(),
-    send: vi.fn(),
-    addEventListener: vi.fn(),
-    removeEventListener: vi.fn(),
+    open: () => current().open(),
+    receive: (message: unknown) => current().receive(message),
+    /** Deliver a notification as the session routes frame it. */
+    send: (notification: unknown) =>
+      current().receive({ event: "marimo", data: { marimo: notification } }),
+    /** End the connection so the stream reconnects. */
+    drop: () => current().drop(),
+    /** Close the connection with a reason the editor treats as final. */
+    close: (reason: string) =>
+      current().stop(
+        new StreamStopped(reason, {
+          cause: classifyCloseEvent({ reason }).status,
+        }),
+      ),
   };
 }
 
 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,22 +147,22 @@ 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()", () => {
-  let transport: MockTransport;
+describe("useNotebookSession.reconnect()", () => {
+  let fake: FakeTransport;
   let isHealthy: Mock<() => Promise<boolean>>;
   let store: ReturnType<typeof createStore>;
 
   beforeEach(() => {
-    transport = makeTransport(WebSocket.CLOSED);
+    fake = withTransport();
     isHealthy = vi.fn<() => Promise<boolean>>().mockResolvedValue(true);
     store = createStore();
     store.set(connectionAtom, {
@@ -122,59 +170,53 @@ 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(useRuntimeManager).mockReturnValue(makeRuntimeManager(isHealthy));
   });
 
-  function renderUseHook() {
-    return renderConnectionHook(store);
+  async function renderStopped() {
+    const hook = renderConnectionHook(store);
+    await started();
+    await act(async () => connection(fake).close("MARIMO_SHUTDOWN"));
+    return hook;
   }
 
-  it("is a no-op when the transport is already OPEN", async () => {
-    transport.readyState = WebSocket.OPEN;
-    const { result } = renderUseHook();
-    await act(async () => {
-      await result.current.reconnect();
-    });
-    expect(isHealthy).not.toHaveBeenCalled();
-    expect(transport.reconnect).not.toHaveBeenCalled();
-  });
-
-  it("is a no-op when the transport is already CONNECTING", async () => {
-    transport.readyState = WebSocket.CONNECTING;
-    const { result } = renderUseHook();
-    await act(async () => {
-      await result.current.reconnect();
-    });
-    expect(isHealthy).not.toHaveBeenCalled();
-    expect(transport.reconnect).not.toHaveBeenCalled();
-  });
+  it.each(["connecting", "connected"])(
+    "is a no-op while the stream is %s",
+    async (state) => {
+      const { result } = renderConnectionHook(store);
+      await started();
+      if (state === "connected") {
+        act(() => connection(fake).open());
+      }
+      await act(async () => {
+        await result.current.reconnect();
+      });
+      expect(isHealthy).not.toHaveBeenCalled();
+      expect(fake.attempts).toHaveLength(1);
+    },
+  );
 
   it("probes /health and reconnects when the runtime is healthy", async () => {
     isHealthy.mockResolvedValue(true);
-    const { result } = renderUseHook();
+    const { result } = await renderStopped();
     await act(async () => {
       await result.current.reconnect();
     });
     expect(isHealthy).toHaveBeenCalledOnce();
-    expect(transport.reconnect).toHaveBeenCalledOnce();
+    expect(fake.attempts).toHaveLength(2);
     expect(store.get(connectionAtom)).toEqual({
       state: ConnectionState.CONNECTING,
     });
   });
 
-  it("transitions to CLOSED and does not call ws.reconnect when the probe fails", async () => {
+  it("transitions to CLOSED and does not reconnect when the probe fails", async () => {
     isHealthy.mockResolvedValue(false);
-    const { result } = renderUseHook();
+    const { result } = await renderStopped();
     await act(async () => {
       await result.current.reconnect();
     });
     expect(isHealthy).toHaveBeenCalledOnce();
-    expect(transport.reconnect).not.toHaveBeenCalled();
+    expect(fake.attempts).toHaveLength(1);
     expect(store.get(connectionAtom)).toEqual({
       state: ConnectionState.CLOSED,
       code: WebSocketClosedReason.KERNEL_DISCONNECTED,
@@ -183,30 +225,31 @@ describe("useMarimoKernelConnection.reconnect()", () => {
   });
 });
 
-describe("useMarimoKernelConnection messages", () => {
-  it("records completion of the initial run", () => {
-    const store = createStore();
-    vi.mocked(useConnectionTransport).mockClear();
-    vi.mocked(useConnectionTransport).mockReturnValue(
-      makeTransport(WebSocket.OPEN),
-    );
-    vi.mocked(useRuntimeManager).mockReturnValue(
-      makeRuntimeManager() as unknown as ReturnType<typeof useRuntimeManager>,
-    );
+describe("useNotebookSession messages", () => {
+  let store: ReturnType<typeof createStore>;
+  let fake: FakeTransport;
+
+  beforeEach(async () => {
+    store = createStore();
+    fake = withTransport();
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
     renderConnectionHook(store);
+    await started();
+  });
 
-    const options = vi.mocked(useConnectionTransport).mock.calls.at(-1)?.[0];
+  function receive(payload: unknown) {
+    act(() => connection(fake).receive(payload));
+  }
+
+  it("records completion of the initial run", () => {
+    const notification = {
+      op: "completed-run",
+      run_id: null,
+      outcome: "succeeded",
+      error: null,
+    } satisfies NotificationPayload["data"];
     expect(store.get(initialRunCompletedAtom)).toBe(false);
-    act(() => {
-      options?.onMessage(
-        new MessageEvent("message", {
-          data: JSON.stringify({
-            op: "completed-run",
-            data: { op: "completed-run", run_id: null },
-          }),
-        }),
-      );
-    });
+    receive({ id: "42", event: "marimo", data: { marimo: notification } });
     expect(store.get(initialRunCompletedAtom)).toBe(true);
   });
 });
@@ -216,47 +259,34 @@ it.each(["kernel-ready", "reconnected"])(
   async (op) => {
     const store = createStore();
     store.set(connectionAtom, { state: ConnectionState.CONNECTING });
-    vi.mocked(useConnectionTransport).mockClear();
-    vi.mocked(useConnectionTransport).mockReturnValue(
-      makeTransport(WebSocket.OPEN),
-    );
-    vi.mocked(useRuntimeManager).mockReturnValue(
-      makeRuntimeManager() as unknown as ReturnType<typeof useRuntimeManager>,
-    );
+    const fake = withTransport();
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
     renderConnectionHook(store);
-    const options = vi.mocked(useConnectionTransport).mock.calls.at(-1)![0];
+    await started();
+    const conn = connection(fake);
 
-    await act(async () => {
-      await options.onOpen(new Event("open"));
-    });
+    await act(async () => conn.open());
     expect(store.get(connectionAtom).state).toBe(ConnectionState.CONNECTING);
 
-    act(() => {
-      options.onMessage(
-        new MessageEvent("message", {
-          data: JSON.stringify({
-            op,
-            data: {
-              op,
-              cell_ids: [],
-              codes: [],
-              names: [],
-              configs: [],
-              layout: null,
-              resumed: true,
-              ui_values: {},
-              last_executed_code: {},
-              last_execution_time: {},
-              app_config: { width: "normal" },
-              kiosk: false,
-              capabilities: { terminal: false },
-              auto_instantiated: false,
-              consumer_capabilities: { edit: true, interact: true },
-            },
-          }),
-        }),
-      );
-    });
+    act(() =>
+      conn.send({
+        op,
+        cell_ids: [],
+        codes: [],
+        names: [],
+        configs: [],
+        layout: null,
+        resumed: true,
+        ui_values: {},
+        last_executed_code: {},
+        last_execution_time: {},
+        app_config: { width: "normal" },
+        kiosk: false,
+        capabilities: { terminal: false },
+        auto_instantiated: false,
+        consumer_capabilities: { edit: true, interact: true },
+      }),
+    );
     expect(store.get(connectionAtom).state).toBe(ConnectionState.OPEN);
   },
 );
@@ -265,19 +295,16 @@ describe("connection notice", () => {
   beforeEach(() => vi.useFakeTimers());
   afterEach(() => vi.useRealTimers());
 
-  function renderNotice() {
+  async function renderNotice() {
     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>,
-    );
+    const fake = withTransport();
+    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,20 +320,17 @@ describe("connection notice", () => {
         </ErrorBoundary>
       </JotaiProvider>,
     );
-    const options = vi.mocked(useConnectionTransport).mock.calls.at(-1)![0];
+    await started();
+    const conn = connection(fake);
+    // Messages arrive on an open connection.
+    act(() => conn.open());
     const send = (data: NotificationPayload["data"]) =>
-      act(() => {
-        options.onMessage(
-          new MessageEvent("message", {
-            data: JSON.stringify({ op: data.op, data }),
-          }),
-        );
-      });
-    return { store, transport, options, send };
+      act(() => conn.send(data));
+    return { store, fake, conn, send };
   }
 
-  it("keeps elapsed time across startup phases and leaves reconnection to the footer once cells are available", () => {
-    const { store, options, send } = renderNotice();
+  it("keeps elapsed time across startup phases and leaves reconnection to the footer once cells are available", async () => {
+    const { store, conn, send } = await renderNotice();
     send({
       op: "startup-progress",
       phase: "preparing-environment",
@@ -346,7 +370,7 @@ describe("connection notice", () => {
     expect(
       screen.queryByRole("region", { name: "Notebook startup" }),
     ).not.toBeInTheDocument();
-    act(() => options.onClose(new CloseEvent("close")));
+    await act(async () => conn.drop());
     act(() => vi.advanceTimersByTime(500));
     expect(store.get(connectionAtom)).toEqual({
       state: ConnectionState.CONNECTING,
@@ -357,8 +381,8 @@ describe("connection notice", () => {
     ).not.toBeInTheDocument();
   });
 
-  it("restores startup output before appending live chunks and resets it for a new attempt", () => {
-    const { store, send } = renderNotice();
+  it("restores startup output before appending live chunks and resets it for a new attempt", async () => {
+    const { store, send } = await renderNotice();
     const snapshot = {
       op: "startup-progress",
       phase: "starting-kernel",
@@ -384,7 +408,7 @@ describe("connection notice", () => {
   });
 
   it("restores completed startup output without putting a connected kernel back into startup", async () => {
-    const { store, options, send } = renderNotice();
+    const { store, conn, send } = await renderNotice();
     send({ op: "reconnected" });
     send({
       op: "startup-progress",
@@ -402,13 +426,15 @@ describe("connection notice", () => {
       screen.queryByRole("region", { name: "Notebook startup" }),
     ).not.toBeInTheDocument();
 
-    // A new transport starts from the next session's authoritative snapshot.
-    await act(async () => options.onOpen(new Event("open")));
+    // A new connection starts from the next session's authoritative snapshot.
+    await act(async () => conn.drop());
+    await act(async () => vi.advanceTimersByTime(2000));
+    await act(async () => conn.open());
     expect(store.get(startupProgressAtom)).toBeNull();
   });
 
-  it("does not flash a notice when startup finishes within the delay", () => {
-    const { send } = renderNotice();
+  it("does not flash a notice when startup finishes within the delay", async () => {
+    const { send } = await renderNotice();
     send({
       op: "startup-progress",
       phase: "preparing-environment",
@@ -436,16 +462,11 @@ describe("connection notice", () => {
   it.each(["preparing-environment", "starting-kernel"] as const)(
     "keeps a %s failure available until the user retries",
     async (phase) => {
-      const { store, transport, options, send } = renderNotice();
+      const { store, fake, conn, send } = await renderNotice();
       send({ op: "startup-progress", phase, logs: "", log_mode: "replace" });
       const error = "A full diagnostic\nwith <stderr> details";
       send({ op: "kernel-startup-error", error });
-      transport.readyState = WebSocket.CLOSED;
-      act(() =>
-        options.onClose(
-          new CloseEvent("close", { reason: "MARIMO_KERNEL_STARTUP_ERROR" }),
-        ),
-      );
+      await act(async () => conn.close("MARIMO_KERNEL_STARTUP_ERROR"));
       expect(screen.getByRole("status")).toHaveTextContent(
         "Kernel failed to start",
       );
@@ -456,11 +477,11 @@ describe("connection notice", () => {
       expect(screen.getByLabelText("Error details")).toHaveTextContent(error, {
         normalizeWhitespace: false,
       });
-      expect(transport.reconnect).not.toHaveBeenCalled();
+      expect(fake.attempts).toHaveLength(1);
       await act(async () =>
         fireEvent.click(screen.getByRole("button", { name: "Try again" })),
       );
-      expect(transport.reconnect).toHaveBeenCalledOnce();
+      expect(fake.attempts).toHaveLength(2);
       expect(store.get(kernelStartupErrorAtom)).toBeNull();
       expect(
         screen.queryByRole("button", { name: "Try again" }),
@@ -480,25 +501,14 @@ describe("connection notice", () => {
   );
 });
 
-it("replaces environment state from snapshots, then continues live progress", () => {
+it("replaces environment state from snapshots, then continues live progress", async () => {
   const store = createStore();
-  vi.mocked(useConnectionTransport).mockClear();
-  vi.mocked(useConnectionTransport).mockReturnValue(
-    makeTransport(WebSocket.OPEN),
-  );
-  vi.mocked(useRuntimeManager).mockReturnValue(
-    makeRuntimeManager() as unknown as ReturnType<typeof useRuntimeManager>,
-  );
+  const fake = withTransport();
+  vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
   renderConnectionHook(store);
-  const options = vi.mocked(useConnectionTransport).mock.calls.at(-1)?.[0];
+  await started();
   function receive(data: NotificationPayload["data"]) {
-    act(() =>
-      options?.onMessage(
-        new MessageEvent("message", {
-          data: JSON.stringify({ op: data.op, data }),
-        }),
-      ),
-    );
+    act(() => connection(fake).send(data));
   }
   receive({
     op: "environment-operation",

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

@@ -0,0 +1,60 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { CloseDecision } from "../host/connection";
+import { Logger } from "@/utils/Logger";
+import { ConnectionState, WebSocketClosedReason } from "./types";
+
+// Keep server reasons in sync with ws_endpoint.py and ws/*.py.
+// SSE carries the same reasons in close events. Unclassified reasons retry.
+export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
+  switch (event.reason) {
+    case "MARIMO_NO_FILE_KEY":
+    case "MARIMO_NO_SESSION_ID":
+    case "MARIMO_NO_SESSION":
+    case "MARIMO_SHUTDOWN":
+      return {
+        kind: "terminal",
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+          reason: "kernel not found",
+        },
+      };
+    case "MARIMO_UNAUTHORIZED":
+      return {
+        kind: "terminal",
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+          reason: "not authorized",
+        },
+      };
+    case "MARIMO_KIOSK_NOT_ALLOWED":
+      return {
+        kind: "terminal",
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+          reason: "kiosk mode is not available for this session",
+        },
+      };
+    case "MARIMO_KERNEL_STARTUP_ERROR":
+      return {
+        kind: "terminal",
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
+          reason: "Failed to start kernel sandbox",
+        },
+      };
+    default:
+      if (event.reason) {
+        Logger.warn("Unrecognized connection close reason", event.reason);
+      }
+  }
+
+  return {
+    kind: "retry",
+    status: { state: ConnectionState.CONNECTING },
+  };
+}

package/src/core/websocket/state.ts

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