@marimo-team/frontend 0.25.3-dev2 → 0.25.3-dev20

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

@@ -0,0 +1,101 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "../protocol";
+import type { Resource } from "./resource";
+
+/**
+ * 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.
+ */
+export 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> = {
+  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/host/resources/resource.ts

@@ -0,0 +1,59 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "../protocol";
+
+/** An event name from the protocol. */
+export type EventName = Protocol.EventOutput["event"];
+
+/** The decoded event with the given name. */
+export type EventOf<Name extends EventName> = Extract<
+  Protocol.EventOutput,
+  { event: Name }
+>;
+
+/**
+ * Events the stream handles itself.
+ *
+ * They identify the notebook, discard state, or carry a native
+ * notification. Resources do not receive them.
+ */
+export type StreamEventName =
+  | "reset"
+  | "notebook"
+  | "notebook.removed"
+  | "marimo";
+
+/** Events a resource may handle. `ready` marks the end of the snapshot. */
+export type ResourceEventName = Exclude<EventName, StreamEventName>;
+
+/** `snapshot` until `ready` arrives and `live` after it. */
+export type StreamPhase = "snapshot" | "live";
+
+/** What a handler may know about the stream beyond the event. */
+export interface HandlerContext {
+  readonly phase: StreamPhase;
+}
+
+/**
+ * A pure handler for each event the resource reads.
+ *
+ * A handler returns `state` when the event does not apply to it.
+ */
+export type EventHandlers<State> = {
+  readonly [Name in ResourceEventName]?: (
+    state: State,
+    event: EventOf<Name>,
+    context: HandlerContext,
+  ) => State;
+};
+
+/**
+ * One kind of object assembled from its events.
+ *
+ * The stream returns a resource to `initial` on `reset` and when the
+ * notebook is removed.
+ */
+export interface Resource<State> {
+  readonly initial: State;
+  readonly on: EventHandlers<State>;
+}

package/src/core/host/stream.ts

@@ -0,0 +1,213 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { NotebookEvent } from "./events";
+import type * as Protocol from "./protocol";
+import type {
+  EventOf,
+  HandlerContext,
+  Resource,
+  ResourceEventName,
+  StreamPhase,
+} from "./resources/resource";
+
+/** A resource for each entry of a state map under the same name. */
+export 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.
+ */
+export interface StreamState<States> {
+  readonly phase: StreamPhase;
+  /**
+   * 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. */
+export 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: EventOf<ResourceEventName>,
+): 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: EventOf<ResourceEventName>,
+  ): 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/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/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: "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,