@marimo-team/frontend 0.25.2 → 0.25.3-dev12

package/src/core/islands/__tests__/bootstrap.test.ts

@@ -3,6 +3,8 @@
 import { createStore } from "jotai";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { RuntimeState } from "@/core/kernel/RuntimeState";
+import { attachmentsAtom } from "@/core/attachments/state";
+import { Logger } from "@/utils/Logger";
 import { initializeIslands } from "../bootstrap";
 import type { IslandsPyodideBridge } from "../bridge";
 import { islandsHydratedAtom, islandsPendingInitialRunsAtom } from "../state";
@@ -48,6 +50,30 @@ describe("initializeIslands", () => {
     expect(store.get(islandsHydratedAtom)).toBe(false);
     store.set(islandsPendingInitialRunsAtom, new Set([1, 2]));
 
+    const errors = vi.spyOn(Logger, "error");
+    consumeMessage?.({
+      sessionGeneration: 1,
+      message: JSON.stringify({
+        op: "attachments",
+        data: {
+          op: "attachments",
+          attachments: [
+            {
+              id: "a1",
+              kind: "agent",
+              name: "Pi",
+              since: "2026-10-14T00:00:00Z",
+              notebook_id: "nb-test",
+              identity: "pi-thread",
+            },
+          ],
+        },
+      }) as IslandsKernelMessage["message"],
+    });
+    expect(errors).not.toHaveBeenCalled();
+    expect(store.get(attachmentsAtom)).toEqual([]);
+    expect(store.get(islandsPendingInitialRunsAtom)).toEqual(new Set([1, 2]));
+
     consumeMessage?.(completedRun(1));
     expect(store.get(islandsHydratedAtom)).toBe(false);
 

package/src/core/islands/bootstrap.ts

@@ -209,6 +209,7 @@ function handleMessage(
       case "cache-info":
       case "kernel-startup-error":
       case "notebook-document-transaction":
+      case "attachments":
         return;
 
       case "completed-run":

package/src/core/run-app.tsx

@@ -22,7 +22,7 @@ import { getSessionId } from "./kernel/session";
 import { connectionAtom } from "./network/connection";
 import { useRequestClient } from "./network/requests";
 import { isAppConnecting } from "./websocket/connection-utils";
-import { useMarimoKernelConnection } from "./websocket/useMarimoKernelConnection";
+import { useNotebookSession } from "./websocket/useNotebookSession";
 
 interface AppProps {
   appConfig: AppConfig;
@@ -51,7 +51,7 @@ export const RunApp: React.FC<AppProps> = ({ appConfig, hideHeader }) => {
     };
   }, []);
 
-  const { connection, reconnect } = useMarimoKernelConnection({
+  const { connection, reconnect } = useNotebookSession({
     autoInstantiate: true,
     setCells: setCells,
     sessionId: getSessionId(),

package/src/core/runtime/__tests__/runtime.test.ts

@@ -2,14 +2,19 @@
 
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { Mocks } from "@/__mocks__/common";
-import type { SessionId } from "@/core/kernel/session";
+import {
+  generateSessionId,
+  isSessionId,
+  type SessionId,
+} from "@/core/kernel/session";
 import { Logger } from "@/utils/Logger";
 import { RuntimeManager } from "../runtime";
 import type { RuntimeConfig } from "../types";
 
 // Mock the session module
-vi.mock("@/core/kernel/session", () => ({
-  getSessionId: () => "test-session-id" as SessionId,
+vi.mock("@/core/kernel/session", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/core/kernel/session")>()),
+  getSessionId: () => "test-session-id",
 }));
 
 vi.mock("@/utils/Logger", () => ({ Logger: Mocks.quietLogger() }));
@@ -51,6 +56,38 @@ describe("RuntimeManager", () => {
   });
 
   describe("getWsURL", () => {
+    it("uses host notebook events with the existing session and auth parameters", () => {
+      const original = window.location.href;
+      window.history.replaceState({}, "", "/?types=run");
+      try {
+        const runtime = new RuntimeManager({
+          ...mockConfig,
+          url: "https://example.com/nested?file=notebook.py",
+          notebookId: "nb_123",
+        });
+        const sessionId = "s_123456";
+        if (!isSessionId(sessionId)) {
+          throw new Error("Invalid test session ID");
+        }
+        const url = runtime.getWsURL(sessionId);
+        expect({
+          endpoint: `${url.origin}${url.pathname}`,
+          params: Object.fromEntries(url.searchParams),
+        }).toEqual({
+          endpoint:
+            "wss://example.com/nested/api/marimo/v1/notebooks/nb_123/events/ws",
+          params: {
+            file: "notebook.py",
+            types: "marimo",
+            session_id: sessionId,
+            access_token: "test-token",
+          },
+        });
+      } finally {
+        window.history.replaceState({}, "", original);
+      }
+    });
+
     it("should return WebSocket URL with session ID for https", () => {
       const runtime = new RuntimeManager(mockConfig);
       const sessionId = "1234" as SessionId;
@@ -59,6 +96,7 @@ describe("RuntimeManager", () => {
       expect(url.protocol).toBe("wss:");
       expect(url.hostname).toBe("example.com");
       expect(url.pathname).toBe("/ws");
+      expect(url.searchParams.has("types")).toBe(false);
       expect(url.searchParams.get("session_id")).toBe(sessionId);
     });
 
@@ -219,14 +257,20 @@ describe("RuntimeManager", () => {
   });
 
   describe("getSseURL", () => {
-    it("should return an http(s) URL with the session ID", () => {
-      const runtime = new RuntimeManager(mockConfig);
-      const url = runtime.getSseURL("1234" as SessionId);
-
-      expect(url.protocol).toBe("https:");
-      expect(url.pathname).toBe("/sse");
-      expect(url.searchParams.get("session_id")).toBe("1234");
-    });
+    it.each([undefined, "nb_123"])(
+      "uses the legacy session stream when notebookId is %s",
+      (notebookId) => {
+        const runtime = new RuntimeManager({
+          ...mockConfig,
+          url: "https://example.com/nested?file=notebook.py",
+          notebookId,
+        });
+        const sessionId = generateSessionId();
+        expect(runtime.getSseURL(sessionId).toString()).toBe(
+          `https://example.com/nested/sse?file=notebook.py&session_id=${sessionId}`,
+        );
+      },
+    );
 
     it("should never put the auth token in the URL, even cross-origin", () => {
       // SSE requests can send headers, so auth travels in the

package/src/core/runtime/runtime.ts

@@ -9,6 +9,8 @@ import { isStaticNotebook } from "../static/static-state";
 import { isWasm } from "../wasm/utils";
 import type { RuntimeConfig } from "./types";
 
+const HOST_EVENT_TYPES = ["marimo"];
+
 export class RuntimeManager {
   private initialHealthyCheck = new Deferred<void>();
   private config: RuntimeConfig;
@@ -154,7 +156,18 @@ export class RuntimeManager {
    * The WebSocket URL of the runtime.
    */
   public getWsURL(sessionId: SessionId): URL {
-    return this.formatWsURL("/ws", this.getSessionSearchParams(sessionId));
+    const searchParams = this.getSessionSearchParams(sessionId);
+    // TODO(MO-8039): Rediscover hosts and resolve notebooks on reconnect.
+    const notebookId = this.config.notebookId;
+    const path = notebookId
+      ? `/api/marimo/v1/notebooks/${encodeURIComponent(notebookId)}/events/ws`
+      : "/ws";
+    const url = this.formatWsURL(path, searchParams);
+    if (notebookId) {
+      // TODO(MO-8034): Separate host options from notebook query params.
+      url.searchParams.set("types", HOST_EVENT_TYPES.join(","));
+    }
+    return url;
   }
 
   /**

package/src/core/runtime/types.ts

@@ -1,6 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 export interface RuntimeConfig {
+  /** The host notebook for WebSocket events. SSE uses the legacy session stream. */
+  readonly notebookId?: string;
   /**
    * The URL of the runtime server.
    */

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