@marimo-team/frontend 0.25.3-dev4 → 0.25.3-dev40

package/src/core/host/stream.ts

@@ -0,0 +1,211 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { NotebookEvent } from "./events";
+import type * as Protocol from "./protocol";
+import type {
+  HandlerContext,
+  Resource,
+  ResourceEvent,
+} from "./resources/resource";
+
+/** A resource for each entry of a state map under the same name. */
+type Resources<States> = {
+  readonly [Name in keyof States]: Resource<States[Name]>;
+};
+
+/**
+ * The state of one notebook stream.
+ *
+ * Holds the phase, the notebook identity, the resume cursor, and the state
+ * of each resource.
+ */
+interface StreamState<States> {
+  readonly phase: HandlerContext["phase"];
+  /**
+   * The notebook the stream describes.
+   *
+   * Retained across `reset`.
+   */
+  readonly notebookId: Protocol.NotebookIdOutput | null;
+  /**
+   * The id of the last live event.
+   *
+   * Null before `ready` and after `reset`.
+   */
+  readonly cursor: string | null;
+  /** The state of each resource by name. */
+  readonly resources: States;
+}
+
+/** A reducer over the events of one notebook stream. */
+interface Stream<States> {
+  readonly initial: StreamState<States>;
+  /**
+   * Apply one event.
+   *
+   * Throws on a live event without a cursor or on an object from another
+   * notebook.
+   */
+  reduce(state: StreamState<States>, event: NotebookEvent): StreamState<States>;
+  /**
+   * The state to hold while reconnecting.
+   *
+   * A live stream keeps its state and resumes from `cursor`. An incomplete
+   * snapshot starts over.
+   */
+  resume(state: StreamState<States>): StreamState<States>;
+}
+
+/** Events that may arrive without a cursor after `ready`. */
+const UNCURSORED = new Set<NotebookEvent["event"]>([
+  "reset",
+  "marimo",
+  "unknown",
+]);
+
+/**
+ * The notebook the event's object belongs to.
+ *
+ * Undefined for objects that do not name one.
+ */
+function notebookIdOf(
+  event: ResourceEvent,
+): Protocol.NotebookIdOutput | undefined {
+  switch (event.event) {
+    case "document":
+      return event.data.document.notebook_id;
+    case "cell":
+      return event.data.cell.notebook_id;
+    case "analysis":
+      return event.data.analysis.notebook_id;
+    case "runtime":
+      return event.data.runtime.notebook_id;
+    case "attachment":
+      return event.data.attachment.notebook_id;
+    case "run":
+      return event.data.run.notebook_id;
+    default:
+      return undefined;
+  }
+}
+
+/** Compose resources into one stream. */
+export function createStream<States extends object>(
+  resources: Resources<States>,
+): Stream<States> {
+  const names = Object.keys(resources) as (keyof States & string)[];
+  const initialResources = Object.fromEntries(
+    names.map((name) => [name, resources[name].initial]),
+  ) as States;
+  const initial: StreamState<States> = {
+    phase: "snapshot",
+    notebookId: null,
+    cursor: null,
+    resources: initialResources,
+  };
+
+  function identify(
+    state: StreamState<States>,
+    notebookId: Protocol.NotebookIdOutput,
+  ): StreamState<States> {
+    if (state.notebookId === notebookId) {
+      return state;
+    }
+    if (state.notebookId !== null) {
+      throw new Error("Notebook stream changed identity");
+    }
+    return { ...state, notebookId };
+  }
+
+  function clear(state: StreamState<States>): StreamState<States> {
+    return state.resources === initialResources
+      ? state
+      : { ...state, resources: initialResources };
+  }
+
+  function fanOut(
+    state: StreamState<States>,
+    event: ResourceEvent,
+  ): StreamState<States> {
+    const context: HandlerContext = { phase: state.phase };
+    let next = state.resources;
+    for (const name of names) {
+      type State = States[typeof name];
+      // Handlers are keyed by the event name they accept.
+      const handler = resources[name].on[event.event] as
+        | ((
+            state: State,
+            event: Protocol.EventOutput,
+            context: HandlerContext,
+          ) => State)
+        | undefined;
+      if (handler === undefined) {
+        continue;
+      }
+      const before = next[name];
+      const after = handler(before, event, context);
+      if (after !== before) {
+        next = { ...next, [name]: after };
+      }
+    }
+    return next === state.resources ? state : { ...state, resources: next };
+  }
+
+  /** Record the position of a live event. */
+  function advance(
+    state: StreamState<States>,
+    id: string | undefined,
+  ): StreamState<States> {
+    return state.phase === "live" && id !== undefined && id !== state.cursor
+      ? { ...state, cursor: id }
+      : state;
+  }
+
+  return {
+    initial,
+    reduce(state, event) {
+      if (
+        state.phase === "live" &&
+        event.id === undefined &&
+        !UNCURSORED.has(event.event)
+      ) {
+        throw new Error("Live notebook event is missing its cursor");
+      }
+      switch (event.event) {
+        case "marimo":
+        case "unknown":
+          return advance(state, event.id);
+        case "reset":
+          return { ...initial, notebookId: state.notebookId };
+        case "ready":
+          // Handlers see the phase that ended, so a snapshot's ready differs
+          // from one that ends a replay.
+          return { ...fanOut(state, event), phase: "live", cursor: event.id };
+        case "notebook":
+          return advance(identify(state, event.data.notebook.id), event.id);
+        case "notebook.removed":
+          return advance(
+            state.notebookId !== null && event.data.id !== state.notebookId
+              ? state
+              : clear(state),
+            event.id,
+          );
+        default: {
+          const notebookId = notebookIdOf(event);
+          return advance(
+            fanOut(
+              notebookId === undefined ? state : identify(state, notebookId),
+              event,
+            ),
+            event.id,
+          );
+        }
+      }
+    },
+    resume(state) {
+      return state.phase === "snapshot"
+        ? { ...initial, notebookId: state.notebookId }
+        : state;
+    },
+  };
+}

package/src/core/hotkeys/actions.ts

@@ -2,6 +2,7 @@
 import { atom, useAtomValue, useSetAtom } from "jotai";
 import useEvent from "react-use-event-hook";
 import type { HotkeyAction } from "@/core/hotkeys/hotkeys";
+import { Objects } from "@/utils/objects";
 
 /**
  * Map of registered keyboard shortcuts and their callbacks.
@@ -29,10 +30,7 @@ export function useSetRegisteredAction() {
       },
     ),
     unregisterAction: useEvent((shortcut: HotkeyAction) => {
-      set((actions) => {
-        const { [shortcut]: unused, ...rest } = actions;
-        return rest;
-      });
+      set((actions) => Objects.omit(actions, [shortcut]));
     }),
   };
 }

package/src/core/network/connection.ts

@@ -7,8 +7,10 @@ import { waitFor } from "../state/jotai";
 import { type ConnectionStatus, ConnectionState } from "../websocket/types";
 
 /**
- * Atom for storing the connection status.
- * Initialized to NOT_STARTED.
+ * Shared application state, not the actual event-stream status.
+ *
+ * TODO(MO-8051) Expose read-only connection state and derive UI availability
+ * separately.
  */
 export const connectionAtom = atom<ConnectionStatus>({
   state: ConnectionState.NOT_STARTED,

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: "notebook,document,cell,attachment,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

@@ -3,12 +3,15 @@
 import { Deferred } from "@/utils/Deferred";
 import { Logger } from "@/utils/Logger";
 import { KnownQueryParams } from "../constants";
+import { NOTEBOOK_EVENT_TYPES } from "../host/events";
 import { isIslands } from "../islands/utils";
 import { getSessionId, type SessionId } from "../kernel/session";
 import { isStaticNotebook } from "../static/static-state";
 import { isWasm } from "../wasm/utils";
 import type { RuntimeConfig } from "./types";
 
+const knownQueryParams = new Set(Object.values(KnownQueryParams));
+
 export class RuntimeManager {
   private initialHealthyCheck = new Deferred<void>();
   private config: RuntimeConfig;
@@ -71,10 +74,7 @@ export class RuntimeManager {
     }
 
     for (const [key, value] of currentParams.entries()) {
-      if (
-        restrictToKnownQueryParams &&
-        !Object.values(KnownQueryParams).includes(key)
-      ) {
+      if (restrictToKnownQueryParams && !knownQueryParams.has(key)) {
         continue;
       }
       baseUrl.searchParams.set(key, value);
@@ -154,7 +154,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", NOTEBOOK_EVENT_TYPES.join(","));
+    }
+    return url;
   }
 
   /**
@@ -164,6 +175,8 @@ export class RuntimeManager {
    * (see `headers()`) instead of the URL.
    */
   public getSseURL(sessionId: SessionId): URL {
+    // TODO: Migrate editor SSE to the host event stream. Until then, host
+    // state such as attachments is unavailable on this legacy connection.
     const url = this.formatHttpURL({
       path: "/sse",
       searchParams: this.getSessionSearchParams(sessionId),

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,170 @@
+/* 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() {
+  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,
+    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 = {
+    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("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("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();
+  });
+});