@marimo-team/frontend 0.25.3-dev129 → 0.25.3-dev134

package/src/core/host/AGENTS.md

@@ -4,13 +4,9 @@ Production code may import host modules, `@marimo-team/client`,
 `jotai/vanilla`, `zod`, and `@/utils/Logger`. Enforced by
 `marimo/host-imports`. Tests are exempt.
 
+Code outside this directory imports it as `@/core/host`. Enforced by
+`marimo/module-entry`.
+
 The protocol schemas and types live in `@marimo-team/client/protocol`.
 
 Define interfaces here and inject application behavior through adapters.
-
-`notebook-client.ts` follows the notebook with `NotebookEvents` from
-`@marimo-team/client`, over a `NotebookTransport` the editor supplies. The
-editor's transports for its session routes live in
-`core/websocket/transports/notebook.ts`. Each kind of object is a `NotebookResource`
-in `resources/` with pure handlers keyed by event name, composed with
-`createNotebookReducer`. The client holds one state atom and exposes derived atoms.

package/src/core/host/__tests__/fixtures.ts

@@ -2,7 +2,37 @@
 
 import NodeFs from "node:fs";
 import NodeUrl from "node:url";
+import type { NotebookTransport } from "@marimo-team/client";
 import * as Protocol from "@marimo-team/client/protocol";
+import { vi } from "vitest";
+import type { NotebookBackend } from "../backend";
+
+type Commands = ReturnType<NotebookBackend["notebook"]>;
+
+/** A backend command, with the notebook it was sent to. */
+export type FakeEdit = (
+  notebookId: Protocol.NotebookId,
+  ...args: Parameters<Commands["edit"]>
+) => ReturnType<Commands["edit"]>;
+
+/** A backend over `transport` whose edits go to `edit`. */
+export function fakeBackend({
+  transport = async () => undefined,
+  edit = vi.fn<FakeEdit>(),
+  checkHealth,
+}: {
+  transport?: NotebookTransport;
+  edit?: FakeEdit;
+  checkHealth?: () => Promise<boolean>;
+} = {}): NotebookBackend {
+  return {
+    transport,
+    notebook: (notebookId) => ({
+      edit: (...args) => edit(notebookId, ...args),
+    }),
+    checkHealth,
+  };
+}
 
 /** Read the protocol's portable examples as transport messages. */
 export function fixtureMessages(name: string) {

package/src/core/host/__tests__/{notebook-client.test.ts → replica.test.ts}

@@ -3,13 +3,12 @@
 
 import { type NotebookTransport, StreamStopped } from "@marimo-team/client";
 import { createStore } from "jotai";
-import { afterEach, describe, expect, it, vi } from "vitest";
+import { afterEach, describe, expect, it, onTestFinished, vi } from "vitest";
 import { Deferred } from "@/utils/Deferred";
 import { ConnectionState } from "../../websocket/types";
-import type { NotebookBackend } from "../backend";
 import type { ConnectionStatus } from "../connection";
-import { createNotebookClient } from "../notebook-client";
-import { fixtureEvents, fixtureMessages } from "./fixtures";
+import { createNotebookReplica } from "../replica";
+import { fakeBackend, fixtureEvents, fixtureMessages } from "./fixtures";
 
 /** One connection the fake transport opened. */
 interface Attempt {
@@ -23,14 +22,7 @@ interface Attempt {
   stop(status?: ConnectionStatus): void;
 }
 
-type Commands = ReturnType<NotebookBackend["notebook"]>;
-/** A backend edit with the notebook it was sent to. */
-type Edit = (
-  id: string,
-  ...args: Parameters<Commands["edit"]>
-) => ReturnType<Commands["edit"]>;
-
-function setup(edit = vi.fn<Edit>()) {
+function setup() {
   const attempts: Attempt[] = [];
   const transport: NotebookTransport = ({ cursor, signal }, handlers) =>
     new Promise((resolve, reject) => {
@@ -48,19 +40,15 @@ function setup(edit = vi.fn<Edit>()) {
   const onNotification = vi.fn();
   const checkHealth = vi.fn(async () => true);
   const store = createStore();
-  const backend: NotebookBackend = {
-    transport,
-    notebook: (id) => ({ edit: (...args) => edit(id, ...args) }),
-  };
-  const controller = createNotebookClient(store);
-  const client = controller.client;
-  const connect = (signal: AbortSignal) =>
-    controller.connect({
-      backend,
-      signal,
-      checkHealth,
+  const backend = fakeBackend({ transport, checkHealth });
+  const client = createNotebookReplica(store);
+  /** Follow the backend until `signal` aborts. */
+  const connect = (signal: AbortSignal) => {
+    const disconnect = client.connect(backend, {
       onNotification: (notification) => onNotification({ data: notification }),
     });
+    signal.addEventListener("abort", disconnect, { once: true });
+  };
   let next = 0;
   /** Wait for the next connection, and open it. */
   const opened = async () => {
@@ -71,6 +59,7 @@ function setup(edit = vi.fn<Edit>()) {
   };
   return {
     attempts,
+    backend,
     client,
     connect,
     onNotification,
@@ -109,61 +98,40 @@ const notebook = {
   updated_at: null,
 };
 
-describe("NotebookClient lifetime", () => {
-  it("keeps an edit's target and request identity across a stream reset", async () => {
-    const edit = vi.fn<Edit>().mockResolvedValue({
-      ok: false,
-      status: 409,
-      problem: { type: "conflict", title: "conflict", status: 409 },
-    });
-    const { connect, opened, client } = setup(edit);
-    connect(lifetime().signal);
+describe("NotebookReplica lifetime", () => {
+  it("ends the previous connection when it connects again", async () => {
+    const { attempts, backend, client, opened, store } = setup();
+    const first = client.connect(backend, { onNotification: vi.fn() });
     const { receive } = await opened();
     fixtureMessages("notebook-events.sse").forEach(receive);
-    const options = { idempotencyKey: "edit-1", signal: lifetime().signal };
-    await client.edit("n1", [], options);
-    receive({ event: "reset", data: {} });
-    await client.edit("n1", [], options);
-    expect(edit.mock.calls).toEqual([
-      ["n1", [], options],
-      ["n1", [], options],
-    ]);
-
-    receive({
-      event: "notebook",
-      data: { notebook: { ...notebook, id: "n2" } },
+    const second = client.connect(backend, { onNotification: vi.fn() });
+    onTestFinished(second);
+    expect(attempts[0].signal.aborted).toBe(true);
+    expect(store.get(client.documentAtom)).toBeNull();
+    await opened();
+    first();
+    expect(attempts[1].signal.aborted).toBe(false);
+    expect(store.get(client.connectionAtom)).toEqual({
+      state: ConnectionState.OPEN,
     });
-    await expect(client.edit("n1", [], options)).rejects.toThrow(
-      "notebook changed",
-    );
-    expect(edit).toHaveBeenCalledTimes(2);
   });
 
-  it("rejects aborted edits and reports edits before connecting", async () => {
-    const { client } = setup();
-    const controller = lifetime();
-    await expect(
-      client.edit("n1", [], {
-        idempotencyKey: "edit-1",
-        signal: controller.signal,
-      }),
-    ).resolves.toEqual({
-      ok: false,
-      status: 501,
-      problem: {
-        type: "tag:marimo.io,2026:host/unsupported-operation",
-        title: "Document editing is unavailable on this connection",
-        status: 501,
-      },
-    });
-    const reason = new Error("Session ended");
-    controller.abort(reason);
-    await expect(
-      client.edit("n1", [], {
-        idempotencyKey: "edit-1",
-        signal: controller.signal,
-      }),
-    ).rejects.toBe(reason);
+  it("exposes the backend it follows until that connection ends", async () => {
+    const { backend, client, opened, store } = setup();
+    expect(store.get(client.backendAtom)).toBeNull();
+    const disconnect = client.connect(backend, { onNotification: vi.fn() });
+    expect(store.get(client.backendAtom)).toBe(backend);
+    const { receive, drop } = await opened();
+    fixtureMessages("notebook-events.sse").forEach(receive);
+    receive({ event: "reset", data: {} });
+    drop();
+    expect(store.get(client.backendAtom)).toBe(backend);
+
+    const other = fakeBackend();
+    onTestFinished(client.connect(other, { onNotification: vi.fn() }));
+    expect(store.get(client.backendAtom)).toBe(other);
+    disconnect();
+    expect(store.get(client.backendAtom)).toBe(other);
   });
 
   it("exposes identity and capabilities only after the notebook is ready", async () => {
@@ -450,7 +418,7 @@ describe("NotebookClient lifetime", () => {
   });
 });
 
-describe("NotebookClient attachments", () => {
+describe("NotebookReplica attachments", () => {
   const [browser, agent] = fixtureEvents("notebook-events.sse").flatMap(
     (event) => (event.event === "attachment" ? [event.data.attachment] : []),
   );

package/src/core/host/__tests__/wasm.test.ts

@@ -99,6 +99,14 @@ describe("WASM notebook backend", () => {
     ).rejects.toBeInstanceOf(StreamStopped);
   });
 
+  it("refuses commands for any notebook but the worker's", () => {
+    const { bridge } = worker();
+    expect(() => createWasmNotebookBackend(bridge).notebook("n2")).toThrow(
+      "serves notebook n1, not n2",
+    );
+    expect(bridge.call).not.toHaveBeenCalled();
+  });
+
   it("sends edits as calls with the retry key and validates replies", async () => {
     const { bridge, reply } = worker();
     const backend = createWasmNotebookBackend(bridge);

package/src/core/host/backend.ts

@@ -1,5 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+/** A notebook's events and commands, from a host or a stand-in for one. */
+
 import type { NotebookClient, NotebookTransport } from "@marimo-team/client";
 import type { NotebookId } from "@marimo-team/client/protocol";
 
@@ -12,4 +14,6 @@ export interface NotebookBackend {
    * Method presence is not permission. The notebook advertises its operations.
    */
   notebook(id: NotebookId): Pick<NotebookClient, "edit">;
+  /** Whether the host is reachable before an explicit reconnect. Absent means always. */
+  checkHealth?(): Promise<boolean>;
 }

package/src/core/host/index.ts

@@ -0,0 +1,14 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** Notebook backends and the replica that follows them. */
+
+export type { NotebookBackend } from "./backend";
+export {
+  type CloseDecision,
+  ConnectionClosedReason,
+  ConnectionState,
+  type ConnectionStatus,
+} from "./connection";
+export { createNotebookReplica, type NotebookReplica } from "./replica";
+export type { Document } from "./resources/document";
+export { createWasmNotebookBackend } from "./wasm";

package/src/core/host/{notebook-client.ts → replica.ts}

@@ -1,20 +1,15 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+/** What a backend's stream has confirmed about a notebook, kept in atoms. */
+
 import { Logger } from "@/utils/Logger";
 import {
   createNotebookReducer,
   type NotebookEvent,
   NotebookEvents,
-  type Result,
   StreamStopped,
 } from "@marimo-team/client";
-import type {
-  Attachment,
-  Change,
-  EditResponseOutput,
-  NotebookId,
-  NotebookOutput,
-} from "@marimo-team/client/protocol";
+import type { Attachment, NotebookOutput } from "@marimo-team/client/protocol";
 import { type Atom, atom, type createStore } from "jotai/vanilla";
 import { z } from "zod";
 import type { NotebookBackend } from "./backend";
@@ -34,9 +29,9 @@ type EditorEvent = NotebookEvent<(typeof TYPES)[number]>;
 /** A native marimo notification, passed through without interpretation. */
 const MarimoMessage = z.object({ marimo: z.unknown() });
 
-/** A connection to a notebook on a marimo host. */
-export interface NotebookClient {
-  /** Client connection status, independent of application session readiness. */
+/** A notebook as its backend's stream reports it, kept current while connected. */
+export interface NotebookReplica {
+  /** Connection status, independent of application session readiness. */
   readonly connectionAtom: Atom<ConnectionStatus>;
   /** Identity and capabilities, available while the stream is connected and ready. */
   readonly notebookAtom: Atom<NotebookOutput | null>;
@@ -48,30 +43,32 @@ export interface NotebookClient {
    * retains the last document. Partial updates are not exposed.
    */
   readonly documentAtom: Atom<Document | null>;
-  /** Apply an edit to the notebook captured by the caller, even across a reset. */
-  edit(
-    notebookId: NotebookId,
-    changes: readonly Change[],
-    options: { readonly idempotencyKey: string; readonly signal: AbortSignal },
-  ): Promise<Result<EditResponseOutput>>;
+  /**
+   * The backend this replica follows, or null while not connected.
+   *
+   * Its commands reach the host that confirmed `documentAtom`. A temporary
+   * disconnection keeps it. Connecting again or disconnecting replaces it.
+   */
+  readonly backendAtom: Atom<NotebookBackend | 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. */
+  /**
+   * Follow a backend, replacing any current connection.
+   *
+   * The returned function ends this connection and clears the replica.
+   * It does nothing once another connection has replaced this one.
+   */
+  connect(
+    backend: NotebookBackend,
+    options: {
+      /** Handle a native marimo notification. Its payload has not been validated. */
+      onNotification: (notification: unknown) => void;
+    },
+  ): () => void;
+  /** Restore the current connection after its stream closed, if the host is healthy. */
   reconnect(): Promise<void>;
 }
 
-/** The backend and application callbacks needed to follow a notebook. */
-export interface NotebookConnection {
-  /** Follows the notebook's stream and sends its commands. */
-  backend: NotebookBackend;
-  /** Abort to close the stream and clear the client's state. */
-  signal: AbortSignal;
-  /** Check whether the host is reachable before an explicit reconnect. */
-  checkHealth: () => Promise<boolean>;
-  /** Handle a native marimo notification. Its payload has not been validated. */
-  onNotification: (notification: unknown) => void;
-}
-
 const disconnected = {
   state: ConnectionState.CLOSED,
   code: ConnectionClosedReason.KERNEL_DISCONNECTED,
@@ -103,8 +100,10 @@ function statusAfter(error: unknown): ConnectionStatus {
   return disconnected;
 }
 
-/** Create a client without opening a connection. */
-export function createNotebookClient(store: ReturnType<typeof createStore>) {
+/** Create a replica without connecting it. */
+export function createNotebookReplica(
+  store: ReturnType<typeof createStore>,
+): NotebookReplica {
   // `useNotebookSession` copies this into the editor's shared `connectionAtom`,
   // which gates features until MO-8051 derives availability separately.
   const connectionAtom = atom<ConnectionStatus>({
@@ -130,50 +129,40 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
       ? get(streamAtom).resources.attachments
       : attachments.initial,
   );
-  let reconnect: (() => Promise<void>) | undefined;
-  let backend: NotebookBackend | undefined;
-  const client: NotebookClient = {
-    connectionAtom,
-    notebookAtom,
-    documentAtom,
-    edit: async (notebookId, changes, options) => {
-      options.signal.throwIfAborted();
-      const currentId = store.get(streamAtom).notebookId;
-      if (currentId !== null && currentId !== notebookId) {
-        throw new Error("The notebook changed before the edit was sent");
-      }
-      if (backend === undefined) {
-        return {
-          ok: false,
-          status: 501,
-          problem: {
-            type: "tag:marimo.io,2026:host/unsupported-operation",
-            title: "Document editing is unavailable on this connection",
-            status: 501,
-          },
-        };
-      }
-      return backend.notebook(notebookId).edit(changes, options);
-    },
-    attachmentsAtom,
-    reconnect: async () => reconnect?.(),
-  };
+  const backendAtom = atom<NotebookBackend | null>(null);
+  /** The connection this replica follows. Connecting again ends it. */
+  let current: { end(): void; reconnect(): Promise<void> } | undefined;
 
   function setStatus(status: ConnectionStatus) {
     store.set(readyAtom, false);
     store.set(connectionAtom, status);
   }
 
-  function connect({
-    backend: connected,
-    signal,
-    checkHealth,
-    onNotification,
-  }: NotebookConnection) {
-    if (signal.aborted) {
-      return;
-    }
-    backend = connected;
+  /** Drop what the stream reported and show `status`. */
+  function clear(status: ConnectionStatus) {
+    store.set(streamAtom, reducer.initial);
+    setStatus(status);
+  }
+
+  return {
+    connectionAtom,
+    notebookAtom,
+    documentAtom,
+    backendAtom,
+    attachmentsAtom,
+    connect,
+    reconnect: async () => current?.reconnect(),
+  };
+
+  function connect(
+    connected: NotebookBackend,
+    { onNotification }: { onNotification: (notification: unknown) => void },
+  ) {
+    current?.end();
+    const lifetime = new AbortController();
+    const { signal } = lifetime;
+    const checkHealth = connected.checkHealth ?? (async () => true);
+    store.set(backendAtom, connected);
     const { transport } = connected;
     let events: NotebookEvents<(typeof TYPES)[number]> | undefined;
     let checkingHealth = false;
@@ -222,8 +211,7 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
       source.close();
       // The cursor can already include an event the replica rejected.
       resumable = false;
-      store.set(streamAtom, reducer.initial);
-      setStatus(malformed);
+      clear(malformed);
     }
 
     function apply(source: NotebookEvents<string>, event: EditorEvent) {
@@ -250,7 +238,7 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
       }
     }
 
-    reconnect = async () => {
+    const reconnect = async () => {
       if (signal.aborted || checkingHealth || events?.status !== "closed") {
         return;
       }
@@ -276,22 +264,23 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
       }
     };
 
-    signal.addEventListener(
-      "abort",
-      () => {
-        reconnect = undefined;
-        if (backend === connected) {
-          backend = undefined;
+    const connection = {
+      end: () => {
+        if (signal.aborted) {
+          return;
         }
-        store.set(streamAtom, reducer.initial);
-        setStatus({ state: ConnectionState.NOT_STARTED });
+        lifetime.abort();
+        if (current === connection) {
+          current = undefined;
+          store.set(backendAtom, null);
+        }
+        clear({ state: ConnectionState.NOT_STARTED });
       },
-      { once: true },
-    );
-
+      reconnect,
+    };
+    current = connection;
     setStatus({ state: ConnectionState.CONNECTING });
     follow();
+    return connection.end;
   }
-
-  return { client, connect };
 }

package/src/core/host/static.ts

@@ -1,5 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+/** A backend that serves a document snapshot without a host. */
+
 import type * as Protocol from "@marimo-team/client/protocol";
 import type { NotebookBackend } from "./backend";
 

package/src/core/host/wasm.ts

@@ -1,5 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+/** A backend for the notebook running in the WASM worker. */
+
 import {
   NotebookClient,
   StreamStopped,
@@ -88,6 +90,13 @@ export function createWasmNotebookBackend(
         signal.removeEventListener("abort", abort);
       }
     },
-    notebook: () => client,
+    notebook: (id) => {
+      if (id !== bridge.notebookId) {
+        throw new Error(
+          `The worker serves notebook ${bridge.notebookId}, not ${id}`,
+        );
+      }
+      return client;
+    },
   };
 }

package/src/core/wasm/bridge.ts

@@ -11,7 +11,7 @@ import { reloadSafe } from "@/utils/reload-safe";
 import { generateUUID } from "@/utils/uuid";
 import { createModuleWorker } from "@/utils/worker";
 import { notebookIsRunningAtom } from "../cells/cells";
-import { createWasmNotebookBackend } from "../host/wasm";
+import { createWasmNotebookBackend } from "../host";
 import type { ClientCommand } from "../kernel/messages";
 import { getInitialAppMode } from "../mode";
 import { API } from "../network/api";