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

package/src/core/host/notebook-client.ts

@@ -0,0 +1,323 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { Logger } from "@/utils/Logger";
+import type { Attachment } from "@marimo-team/client/protocol";
+import { createReducer } from "@marimo-team/client/reducer";
+import { type Atom, atom, type createStore } from "jotai/vanilla";
+import {
+  type ConnectionStatus,
+  ConnectionState,
+  ConnectionClosedReason,
+  type CloseDecision,
+  type NotebookTransport,
+  TransportReadyState,
+} from "./connection";
+import { type NotebookEvent, parseNotebookEvent } from "./events";
+import { attachments } from "./resources/attachments";
+import { type Document, document } from "./resources/document";
+
+/** Events that may arrive without a cursor after `ready`. */
+const UNCURSORED = new Set<NotebookEvent["event"]>([
+  "reset",
+  "marimo",
+  "unknown",
+]);
+
+/** A connection to a notebook on a marimo host. */
+export interface NotebookClient {
+  /** Client connection status, independent of application session readiness. */
+  readonly connectionAtom: Atom<ConnectionStatus>;
+  /**
+   * The most recently confirmed document.
+   *
+   * Null until a snapshot is complete, after reset or removal, and on
+   * connections without host document events. A temporary disconnection
+   * retains the last document. Partial updates are not exposed.
+   */
+  readonly documentAtom: Atom<Document | null>;
+  /** Live attachments, empty until the stream is ready or while disconnected. */
+  readonly attachmentsAtom: Atom<readonly Attachment[]>;
+  /** Attempt to restore a closed connection after checking host health. */
+  reconnect(): Promise<void>;
+}
+
+/** The transport and application callbacks needed to connect a notebook. */
+export interface NotebookConnection {
+  /** Create a transport that reads the URL again on each connection attempt. */
+  createTransport: (url: () => string) => NotebookTransport;
+  /** The current stream URL, including authentication and subscription options. */
+  url: () => string;
+  /** Host events require an intact stream; legacy notifications are independent. */
+  messageFormat: "host" | "session";
+  /** Abort to close the transport and clear the client's state. */
+  signal: AbortSignal;
+  /** Wait until the host can accept the initial connection. */
+  waitUntilReady: () => Promise<void>;
+  /** Check whether the host is reachable before an explicit reconnect. */
+  checkHealth: () => Promise<boolean>;
+  /** Decode wire text without interpreting its event type. */
+  decodeJson: (message: string) => unknown;
+  /** Handle a native marimo notification. Its payload has not been validated. */
+  onNotification: (notification: unknown) => void;
+  /** Report a malformed message without coupling the client to editor UI. */
+  onMessageError: (error: unknown) => void;
+  /** Choose the connection status and retry policy after the transport closes. */
+  classifyClose: (event: CloseEvent) => CloseDecision;
+}
+
+/** Create a client without opening a connection. */
+export function createNotebookClient(store: ReturnType<typeof createStore>) {
+  const connectionAtom = atom<ConnectionStatus>({
+    state: ConnectionState.NOT_STARTED,
+  });
+  // One atom holds the whole stream, which resource atoms derive from.
+  const reducer = createReducer({ document, attachments });
+  const streamAtom = atom(reducer.initial);
+  // A resumed stream is live, but presence waits for this connection's ready.
+  const readyAtom = atom(false);
+  const documentAtom = atom(
+    (get) => get(streamAtom).resources.document.current,
+  );
+  const attachmentsAtom = atom((get) =>
+    get(readyAtom) && get(connectionAtom).state === ConnectionState.OPEN
+      ? get(streamAtom).resources.attachments
+      : attachments.initial,
+  );
+  let reconnect: (() => Promise<void>) | undefined;
+  const client: NotebookClient = {
+    connectionAtom,
+    documentAtom,
+    attachmentsAtom,
+    reconnect: async () => reconnect?.(),
+  };
+
+  function connect({
+    createTransport,
+    url,
+    messageFormat,
+    signal,
+    waitUntilReady,
+    checkHealth,
+    onNotification,
+    onMessageError,
+    decodeJson,
+    classifyClose,
+  }: NotebookConnection) {
+    if (signal.aborted) {
+      return;
+    }
+    let shouldRetry = true;
+    let terminal = false;
+    let checkingHealth = false;
+    let streamKey: string | undefined;
+    // The id of the last live event. Null before `ready` and after `reset`.
+    let cursor: string | null = null;
+    const transport = createTransport(() => {
+      if (messageFormat === "session") {
+        return url();
+      }
+      const next = new URL(url());
+      next.searchParams.delete("cursor");
+      if (signal.aborted) {
+        return next.toString();
+      }
+      // A cursor is useful with the state that produced it, which a change
+      // of session or notebook invalidates. Authentication may change freely.
+      const identity = new URL(next);
+      identity.searchParams.delete("access_token");
+      identity.searchParams.sort();
+      const key = identity.toString();
+      const current = store.get(streamAtom);
+      // A live stream keeps its state and resumes from the cursor. An
+      // incomplete snapshot starts over.
+      const state =
+        key !== streamKey
+          ? reducer.initial
+          : current.phase === "snapshot"
+            ? { ...reducer.initial, notebookId: current.notebookId }
+            : current;
+      if (state.phase === "snapshot") {
+        cursor = null;
+      }
+      streamKey = key;
+      store.set(readyAtom, false);
+      store.set(streamAtom, state);
+      if (cursor !== null) {
+        next.searchParams.set("cursor", cursor);
+      }
+      return next.toString();
+    });
+    const setStatus = (status: ConnectionStatus) => {
+      store.set(readyAtom, false);
+      store.set(connectionAtom, status);
+    };
+    const disconnected = {
+      state: ConnectionState.CLOSED,
+      code: ConnectionClosedReason.KERNEL_DISCONNECTED,
+      reason: "kernel not found",
+    } satisfies ConnectionStatus;
+
+    function retry(code?: number, reason?: string) {
+      if (!signal.aborted && !terminal && shouldRetry) {
+        shouldRetry = false;
+        transport.reconnect(code, reason);
+      }
+    }
+
+    function onOpen() {
+      if (signal.aborted || terminal) {
+        return;
+      }
+      shouldRetry = true;
+      setStatus({ state: ConnectionState.OPEN });
+    }
+
+    function onClose(event: CloseEvent) {
+      if (signal.aborted || terminal) {
+        return;
+      }
+      Logger.warn("Notebook connection closed", event.code, event.reason);
+      const decision = classifyClose(event);
+      setStatus(decision.status);
+      if (decision.kind === "terminal") {
+        terminal = true;
+        transport.close();
+      } else if (decision.kind === "retry") {
+        retry(event.code, event.reason);
+      } else {
+        terminal = true;
+      }
+    }
+
+    function onError(event: Event) {
+      if (signal.aborted || terminal) {
+        return;
+      }
+      Logger.warn("Notebook connection error", event);
+      setStatus(disconnected);
+      retry();
+    }
+
+    function onMessage(message: MessageEvent<string>) {
+      if (signal.aborted || terminal) {
+        return;
+      }
+      let event;
+      try {
+        event = parseNotebookEvent(decodeJson(message.data));
+        const state = store.get(streamAtom);
+        if (
+          state.phase === "live" &&
+          event.id === undefined &&
+          !UNCURSORED.has(event.event)
+        ) {
+          throw new Error("Live notebook event is missing its cursor");
+        }
+        if (event.event !== "marimo" && event.event !== "unknown") {
+          store.set(streamAtom, reducer.reduce(state, event));
+        }
+        if (event.event === "reset") {
+          cursor = null;
+        } else if (event.event === "ready" || state.phase === "live") {
+          cursor = event.id ?? cursor;
+        }
+        if (event.event === "ready" || event.event === "reset") {
+          store.set(readyAtom, event.event === "ready");
+        }
+      } catch (error) {
+        Logger.error("Failed to decode notebook message", error);
+        if (messageFormat === "session") {
+          onMessageError(error);
+          return;
+        }
+        store.set(streamAtom, reducer.initial);
+        cursor = null;
+        terminal = true;
+        setStatus({
+          ...disconnected,
+          reason: "Failed to decode notebook message",
+        });
+        transport.close();
+        return;
+      }
+      if (event.event === "marimo") {
+        onNotification(event.data.marimo);
+      }
+    }
+
+    reconnect = async () => {
+      if (
+        signal.aborted ||
+        checkingHealth ||
+        transport.readyState === TransportReadyState.OPEN ||
+        transport.readyState === TransportReadyState.CONNECTING
+      ) {
+        return;
+      }
+      checkingHealth = true;
+      setStatus({ state: ConnectionState.CONNECTING });
+      try {
+        const healthy = await checkHealth();
+        if (signal.aborted) {
+          return;
+        }
+        if (!healthy) {
+          shouldRetry = false;
+          setStatus(disconnected);
+          return;
+        }
+        terminal = false;
+        shouldRetry = true;
+        transport.reconnect();
+      } catch (error) {
+        if (!signal.aborted) {
+          Logger.error("Notebook health check failed", error);
+          shouldRetry = false;
+          setStatus(disconnected);
+        }
+      } finally {
+        checkingHealth = false;
+      }
+    };
+
+    signal.addEventListener(
+      "abort",
+      () => {
+        transport.removeEventListener("open", onOpen);
+        transport.removeEventListener("close", onClose);
+        transport.removeEventListener("error", onError);
+        transport.removeEventListener("message", onMessage);
+        transport.close();
+        reconnect = undefined;
+        store.set(streamAtom, reducer.initial);
+        setStatus({ state: ConnectionState.NOT_STARTED });
+      },
+      { once: true },
+    );
+
+    transport.addEventListener("open", onOpen);
+    transport.addEventListener("close", onClose);
+    transport.addEventListener("error", onError);
+    transport.addEventListener("message", onMessage);
+
+    if (transport.readyState === TransportReadyState.CLOSED) {
+      setStatus({ state: ConnectionState.CONNECTING });
+      void waitUntilReady()
+        .then(() => {
+          if (!signal.aborted) {
+            transport.reconnect();
+          }
+        })
+        .catch((error) => {
+          if (!signal.aborted) {
+            Logger.error("Healthy connection never made", error);
+            terminal = true;
+            setStatus(disconnected);
+            transport.close();
+          }
+        });
+    }
+  }
+
+  return { client, connect };
+}

package/src/core/host/resources/attachments.ts

@@ -0,0 +1,19 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { Attachment } from "@marimo-team/client/protocol";
+import type { Resource } from "@marimo-team/client/reducer";
+
+/** Attachment records retained with the stream's resume cursor. */
+export const attachments: Resource<readonly Attachment[], "attachment"> = {
+  initial: [],
+  on: {
+    attachment: (state, { data: { attachment } }) =>
+      state.some((item) => item.id === attachment.id)
+        ? state.map((item) => (item.id === attachment.id ? attachment : item))
+        : [...state, attachment],
+    "attachment.removed": (state, { data }) =>
+      state.some((item) => item.id === data.id)
+        ? state.filter((item) => item.id !== data.id)
+        : state,
+  },
+};

package/src/core/host/resources/document.ts

@@ -0,0 +1,101 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import type { Resource } from "@marimo-team/client/reducer";
+
+/**
+ * The notebook's content at one revision.
+ *
+ * The revision record carries the cell order, app config, etags, and
+ * version. Each cell in its order has a record in `cells`.
+ */
+export interface Document {
+  readonly revision: Protocol.Document;
+  readonly cells: ReadonlyMap<Protocol.CellId, Protocol.Cell>;
+}
+
+/** The records of a revision still being received. */
+interface PendingDocument {
+  /** Null until the `document` event that closes the group arrives. */
+  readonly revision: Protocol.Document | null;
+  readonly cells: ReadonlyMap<Protocol.CellId, Protocol.Cell>;
+}
+
+/**
+ * The confirmed document and the revision being assembled after it.
+ *
+ * A snapshot's records are pending until `ready`. A live revision is pending
+ * from its first cell until its `document` event.
+ */
+interface DocumentState {
+  readonly current: Document | null;
+  readonly pending: PendingDocument | null;
+}
+
+function pendingOf(state: DocumentState): PendingDocument {
+  return (
+    state.pending ?? {
+      revision: state.current?.revision ?? null,
+      cells: state.current?.cells ?? new Map(),
+    }
+  );
+}
+
+/** Confirm a revision whose records are complete. */
+function publish(
+  revision: Protocol.Document,
+  cells: ReadonlyMap<Protocol.CellId, Protocol.Cell>,
+): DocumentState {
+  const order = new Set(revision.order);
+  if (
+    order.size !== revision.order.length ||
+    order.size !== cells.size ||
+    [...order].some((id) => !cells.has(id)) ||
+    [...cells.values()].some((cell) => cell.version > revision.version)
+  ) {
+    throw new Error("Document records do not match its revision and order");
+  }
+  return { current: { revision, cells }, pending: null };
+}
+
+export const document: Resource<DocumentState, "cell" | "document"> = {
+  initial: { current: null, pending: null },
+  on: {
+    cell: (state, { data }) => {
+      const pending = pendingOf(state);
+      return {
+        ...state,
+        pending: {
+          revision: pending.revision,
+          cells: new Map(pending.cells).set(data.cell.id, data.cell),
+        },
+      };
+    },
+    "cell.removed": (state, { data }) => {
+      const pending = pendingOf(state);
+      if (!pending.cells.has(data.id)) {
+        return state;
+      }
+      const cells = new Map(pending.cells);
+      cells.delete(data.id);
+      return { ...state, pending: { ...pending, cells } };
+    },
+    document: (state, { data }, { phase }) => {
+      const { cells } = pendingOf(state);
+      return phase === "snapshot"
+        ? { ...state, pending: { revision: data.document, cells } }
+        : publish(data.document, cells);
+    },
+    ready: (state, _event, { phase }) => {
+      if (state.pending === null) {
+        return state;
+      }
+      if (phase === "live") {
+        throw new Error("Stream ended with an incomplete document update");
+      }
+      return state.pending.revision === null
+        ? { current: null, pending: null }
+        : publish(state.pending.revision, state.pending.cells);
+    },
+  },
+};

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.
    */