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

package/src/core/websocket/__tests__/backend.test.ts

@@ -0,0 +1,111 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { NotebookEvents } from "@marimo-team/client";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { createEditorBackend } from "../backend";
+
+const wasm = vi.hoisted(() => ({ enabled: false, consumer: vi.fn() }));
+vi.mock("../../wasm/utils", () => ({ isWasm: () => wasm.enabled }));
+vi.mock("../../wasm/bridge", () => ({
+  PyodideBridge: { INSTANCE: { attachMessageConsumer: wasm.consumer } },
+}));
+
+afterEach(() => {
+  wasm.enabled = false;
+  wasm.consumer.mockReset();
+  vi.unstubAllGlobals();
+});
+
+describe("createEditorBackend", () => {
+  const options = {
+    static: false,
+    transportType: "websocket" as const,
+    url: () => "ws://localhost/ws",
+    headers: () => ({}),
+    notebookUrl: (id: string) =>
+      new URL(`http://localhost/api/marimo/v1/notebooks/${id}`),
+  };
+
+  it.each([
+    ["a static notebook", { ...options, static: true }, "idle"],
+    ["WASM", options, "wasm"],
+    ["SSE", { ...options, transportType: "sse" as const }, "fetch"],
+    ["the default", options, "socket"],
+  ] as const)("follows %s", async (_name, chosen, expected) => {
+    wasm.enabled = expected === "wasm";
+    const fetch = vi.fn(() => new Promise<Response>(() => undefined));
+    const socket = vi.fn();
+    vi.stubGlobal("fetch", fetch);
+    vi.stubGlobal(
+      "WebSocket",
+      class extends EventTarget {
+        public constructor(url: URL) {
+          super();
+          socket(url);
+        }
+        public close() {}
+      },
+    );
+    const events = new NotebookEvents(createEditorBackend(chosen).transport);
+    await vi.waitFor(() =>
+      expect(
+        {
+          idle:
+            events.status === "connected" && !wasm.consumer.mock.calls.length,
+          wasm: wasm.consumer.mock.calls.length === 1,
+          fetch: fetch.mock.calls.length === 1,
+          socket: socket.mock.calls.length === 1,
+        }[expected],
+      ).toBe(true),
+    );
+    events.close();
+  });
+
+  it("edits a reported notebook on its host with the session's headers", async () => {
+    const fetch = vi.fn(
+      async (_input: RequestInfo | URL, _init?: RequestInit) =>
+        new Response(
+          JSON.stringify({ type: "conflict", title: "Conflict", status: 409 }),
+          {
+            status: 409,
+            headers: { "Content-Type": "application/problem+json" },
+          },
+        ),
+    );
+    vi.stubGlobal("fetch", fetch);
+    const backend = createEditorBackend({
+      ...options,
+      headers: () => ({
+        Authorization: "Bearer editor-token",
+        "Marimo-Server-Token": "server-token",
+      }),
+    });
+    const result = await backend
+      .notebook("n1")
+      .edit([], { idempotencyKey: "edit-1" });
+    expect(result).toMatchObject({ ok: false, status: 409 });
+    const [[input, init]] = fetch.mock.calls;
+    expect(String(input)).toBe(
+      "http://localhost/api/marimo/v1/notebooks/n1/edits",
+    );
+    const headers = new Headers(init?.headers);
+    expect(init?.method).toBe("POST");
+    expect(headers.get("Authorization")).toBe("Bearer editor-token");
+    expect(headers.get("Marimo-Server-Token")).toBe("server-token");
+    expect(headers.get("Idempotency-Key")).toBe("edit-1");
+  });
+
+  it.each([
+    ["a static notebook", { ...options, static: true }, false],
+    ["WASM", options, true],
+  ] as const)("does not edit %s", async (_name, chosen, isWasm) => {
+    wasm.enabled = isWasm;
+    const fetch = vi.fn();
+    vi.stubGlobal("fetch", fetch);
+    const result = await createEditorBackend(chosen)
+      .notebook("n1")
+      .edit([], { idempotencyKey: "edit-1" });
+    expect(result).toMatchObject({ ok: false, status: 501 });
+    expect(fetch).not.toHaveBeenCalled();
+  });
+});

package/src/core/websocket/__tests__/useNotebookClient.hook.test.tsx

@@ -8,13 +8,19 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import type { NotebookTransport } from "@marimo-team/client";
 import { ConnectionState } from "../types";
 import { Deferred } from "@/utils/Deferred";
-import { createNotebookTransport } from "../transports/notebook";
+import type { NotebookBackend } from "../../host/backend";
+import { createEditorBackend } from "../backend";
 import { useNotebookClient } from "../useNotebookClient";
 import { fixtureMessages } from "../../host/__tests__/fixtures";
 import { agentAttachmentsAtom } from "../../attachments/state";
 import { notebookClientAtom } from "../state";
 
-vi.mock("../transports/notebook", () => ({ createNotebookTransport: vi.fn() }));
+vi.mock("../backend", () => ({ createEditorBackend: vi.fn() }));
+
+/** A backend that follows `transport` and cannot edit. */
+function backend(transport: NotebookTransport): NotebookBackend {
+  return { transport, notebook: () => ({ edit: vi.fn() }) };
+}
 
 /** A transport whose connections the test drives. */
 function makeTransport() {
@@ -53,6 +59,7 @@ function setup(
     transportType: "websocket" as const,
     url: () => "ws://localhost/ws",
     headers: () => ({}),
+    notebookUrl: (id: string) => new URL(`http://localhost/notebooks/${id}`),
     waitUntilReady: async () => undefined,
     checkHealth: async () => true,
     onNotification: vi.fn(),
@@ -70,7 +77,7 @@ function setup(
 }
 
 beforeEach(() => {
-  vi.mocked(createNotebookTransport).mockReset();
+  vi.mocked(createEditorBackend).mockReset();
 });
 
 describe("useNotebookClient", () => {
@@ -78,16 +85,15 @@ describe("useNotebookClient", () => {
     "rechecks health after failed startup (healthy: %s)",
     async (healthy) => {
       const fake = makeTransport();
-      vi.mocked(createNotebookTransport).mockImplementation(
-        ({ url }) =>
-          async (attempt, handlers) => {
-            if (typeof url === "function") {
-              await url();
-            }
-            if (!attempt.signal.aborted) {
-              return fake.transport(attempt, handlers);
-            }
-          },
+      vi.mocked(createEditorBackend).mockImplementation(({ url }) =>
+        backend(async (attempt, handlers) => {
+          if (typeof url === "function") {
+            await url();
+          }
+          if (!attempt.signal.aborted) {
+            return fake.transport(attempt, handlers);
+          }
+        }),
       );
       const waitUntilReady = vi
         .fn()
@@ -125,14 +131,13 @@ describe("useNotebookClient", () => {
   it("waits for initial readiness without reviving an unmounted connection", async () => {
     const ready = new Deferred<void>();
     const resolvedUrl = vi.fn();
-    vi.mocked(createNotebookTransport).mockImplementation(
-      ({ url }) =>
-        async () => {
-          if (typeof url === "function") {
-            await url();
-            resolvedUrl();
-          }
-        },
+    vi.mocked(createEditorBackend).mockImplementation(({ url }) =>
+      backend(async () => {
+        if (typeof url === "function") {
+          await url();
+          resolvedUrl();
+        }
+      }),
     );
     const { result, store, unmount } = setup(false, {
       waitUntilReady: () => ready.promise,
@@ -152,7 +157,7 @@ describe("useNotebookClient", () => {
 
   it("exposes host presence to the UI without waiting for kernel-ready", async () => {
     const fake = makeTransport();
-    vi.mocked(createNotebookTransport).mockReturnValue(fake.transport);
+    vi.mocked(createEditorBackend).mockReturnValue(backend(fake.transport));
     const { result, store, unmount } = setup();
     await act(async () => Promise.resolve());
     expect(store.get(notebookClientAtom)).toBe(result.current);
@@ -177,7 +182,7 @@ describe("useNotebookClient", () => {
 
   it("uses current callbacks and URL providers without reconnecting on a rerender", async () => {
     const fake = makeTransport();
-    vi.mocked(createNotebookTransport).mockReturnValue(fake.transport);
+    vi.mocked(createEditorBackend).mockReturnValue(backend(fake.transport));
     const { rerender, options, unmount } = setup();
     await act(async () => Promise.resolve());
     const onNotification = vi.fn();
@@ -187,8 +192,8 @@ describe("useNotebookClient", () => {
     expect(onNotification).toHaveBeenCalledExactlyOnceWith({
       data: { op: "reconnected" },
     });
-    expect(createNotebookTransport).toHaveBeenCalledOnce();
-    const [{ url }] = vi.mocked(createNotebookTransport).mock.calls[0];
+    expect(createEditorBackend).toHaveBeenCalledOnce();
+    const [{ url }] = vi.mocked(createEditorBackend).mock.calls[0];
     expect(String(await (url as () => Promise<string>)())).toBe(
       "ws://redirected/ws",
     );
@@ -197,10 +202,10 @@ describe("useNotebookClient", () => {
 
   it("closes every stream under Strict Mode and only follows the active lifetime", async () => {
     const fakes: ReturnType<typeof makeTransport>[] = [];
-    vi.mocked(createNotebookTransport).mockImplementation(() => {
+    vi.mocked(createEditorBackend).mockImplementation(() => {
       const next = makeTransport();
       fakes.push(next);
-      return next.transport;
+      return backend(next.transport);
     });
     const { result, store, options, unmount } = setup(true);
     await act(async () => Promise.resolve());

package/src/core/websocket/__tests__/useNotebookSession.hook.test.tsx

@@ -21,8 +21,8 @@ import {
   vi,
 } from "vitest";
 
-vi.mock("../transports/notebook", () => ({
-  createNotebookTransport: vi.fn(),
+vi.mock("../backend", () => ({
+  createEditorBackend: vi.fn(),
 }));
 
 vi.mock("@/core/runtime/config", async () => {
@@ -51,7 +51,7 @@ import { generateSessionId } from "../../kernel/session";
 import { RuntimeManager } from "../../runtime/runtime";
 import { WebSocketClosedReason, ConnectionState } from "../types";
 import { classifyCloseEvent } from "../connection-state";
-import { createNotebookTransport } from "../transports/notebook";
+import { createEditorBackend } from "../backend";
 import { useNotebookSession } from "../useNotebookSession";
 import { type NotebookTransport, StreamStopped } from "@marimo-team/client";
 
@@ -85,7 +85,10 @@ type FakeTransport = ReturnType<typeof makeTransport>;
 
 /** Use a fake transport for the next session the hook follows. */
 function withTransport(fake: FakeTransport = makeTransport()) {
-  vi.mocked(createNotebookTransport).mockReturnValue(fake.transport);
+  vi.mocked(createEditorBackend).mockReturnValue({
+    transport: fake.transport,
+    notebook: () => ({ edit: vi.fn() }),
+  });
   return fake;
 }
 

package/src/core/websocket/backend.ts

@@ -0,0 +1,68 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { NotebookClient, type StreamUrl } from "@marimo-team/client";
+import type { NotebookId } from "@marimo-team/client/protocol";
+import type { NotebookBackend } from "../host/backend";
+import { PyodideBridge } from "../wasm/bridge";
+import { isWasm } from "../wasm/utils";
+import {
+  idle,
+  pyodide,
+  sessionSse,
+  websocket,
+  withHeaders,
+} from "./transports/notebook";
+
+/** How the editor follows its notebook. */
+export interface EditorBackendOptions {
+  static: boolean;
+  transportType: "websocket" | "sse";
+  /** The session route, read before each connection. */
+  url: StreamUrl;
+  /** The runtime's headers, read before each request. */
+  headers: () => Record<string, string>;
+  /** The host URL of a notebook the stream reported. */
+  notebookUrl: (notebookId: NotebookId) => URL;
+}
+
+/** Commands for a notebook the editor cannot change through its backend. */
+function uneditable(): ReturnType<NotebookBackend["notebook"]> {
+  return {
+    edit: async () => ({
+      ok: false,
+      status: 501,
+      problem: {
+        type: "tag:marimo.io,2026:host/unsupported-operation",
+        title: "Document editing is unavailable on this connection",
+        status: 501,
+      },
+    }),
+  };
+}
+
+/** Choose the backend for the editor's notebook. */
+export function createEditorBackend(
+  options: EditorBackendOptions,
+): NotebookBackend {
+  if (options.static) {
+    return { transport: idle, notebook: uneditable };
+  }
+  if (isWasm()) {
+    return {
+      transport: pyodide((consumer) =>
+        PyodideBridge.INSTANCE.attachMessageConsumer(consumer),
+      ),
+      notebook: uneditable,
+    };
+  }
+  return {
+    transport:
+      options.transportType === "sse"
+        ? sessionSse(options.url, options.headers)
+        : websocket(options.url),
+    notebook: (id) =>
+      new NotebookClient(options.notebookUrl(id), {
+        fetch: withHeaders(options.headers),
+      }),
+  };
+}

package/src/core/websocket/transports/__tests__/notebook.test.ts

@@ -11,18 +11,7 @@ import { createNotebookClient } from "../../../host/notebook-client";
 import { fixtureMessages } from "../../../host/__tests__/fixtures";
 import { ConnectionState } from "../../../host/connection";
 import { Deferred } from "@/utils/Deferred";
-import {
-  createNotebookTransport,
-  pyodide,
-  sessionSse,
-  websocket,
-} from "../notebook";
-
-const wasm = vi.hoisted(() => ({ enabled: false, consumer: vi.fn() }));
-vi.mock("../../../wasm/utils", () => ({ isWasm: () => wasm.enabled }));
-vi.mock("../../../wasm/bridge", () => ({
-  PyodideBridge: { INSTANCE: { attachMessageConsumer: wasm.consumer } },
-}));
+import { pyodide, sessionSse, websocket } from "../notebook";
 
 const encoder = new TextEncoder();
 
@@ -83,14 +72,14 @@ function network(transportType: "websocket" | "sse", url: StreamUrl) {
   );
   const count = () =>
     transportType === "websocket" ? FakeSocket.opened.length : requests.length;
-  const transport = createNotebookTransport({
-    static: false,
-    transportType,
-    url,
-    headers: () => ({}),
-  });
+  const transport =
+    transportType === "websocket"
+      ? websocket(url)
+      : sessionSse(url, () => ({}));
+  const backend = { transport, notebook: () => ({ edit: vi.fn() }) };
   let index = 0;
   return {
+    backend,
     transport,
     count,
     opened: async () => {
@@ -126,8 +115,6 @@ function network(transportType: "websocket" | "sse", url: StreamUrl) {
 
 afterEach(() => {
   FakeSocket.opened.length = 0;
-  wasm.enabled = false;
-  wasm.consumer.mockReset();
   vi.unstubAllGlobals();
   vi.useRealTimers();
 });
@@ -313,7 +300,7 @@ describe.each(["websocket", "sse"] as const)(
         const lifetime = new AbortController();
         onTestFinished(() => lifetime.abort());
         connect({
-          transport: connection.transport,
+          backend: connection.backend,
           signal: lifetime.signal,
           checkHealth: async () => true,
           onNotification: vi.fn(),
@@ -380,7 +367,7 @@ describe.each(["websocket", "sse"] as const)(
       const lifetime = new AbortController();
       onTestFinished(() => lifetime.abort());
       connect({
-        transport: connection.transport,
+        backend: connection.backend,
         signal: lifetime.signal,
         checkHealth: async () => true,
         onNotification: vi.fn(),
@@ -466,37 +453,3 @@ describe("pyodide", () => {
     events.close();
   });
 });
-
-describe("createNotebookTransport", () => {
-  const options = {
-    static: false,
-    transportType: "websocket" as const,
-    url: () => "ws://localhost/ws",
-    headers: () => ({}),
-  };
-
-  it.each([
-    ["a static notebook", { ...options, static: true }, "idle"],
-    ["WASM", options, "wasm"],
-    ["SSE", { ...options, transportType: "sse" as const }, "fetch"],
-    ["the default", options, "socket"],
-  ] as const)("follows %s", async (_name, chosen, expected) => {
-    wasm.enabled = expected === "wasm";
-    const fetch = vi.fn(() => new Promise<Response>(() => undefined));
-    vi.stubGlobal("fetch", fetch);
-    vi.stubGlobal("WebSocket", FakeSocket);
-    const events = new NotebookEvents(createNotebookTransport(chosen));
-    await vi.waitFor(() =>
-      expect(
-        {
-          idle:
-            events.status === "connected" && !wasm.consumer.mock.calls.length,
-          wasm: wasm.consumer.mock.calls.length === 1,
-          fetch: fetch.mock.calls.length === 1,
-          socket: FakeSocket.opened.length === 1,
-        }[expected],
-      ).toBe(true),
-    );
-    events.close();
-  });
-});

package/src/core/websocket/transports/notebook.ts

@@ -7,8 +7,6 @@ import {
   type StreamUrl,
 } from "@marimo-team/client";
 import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
-import { PyodideBridge } from "../../wasm/bridge";
-import { isWasm } from "../../wasm/utils";
 import { classifyCloseEvent } from "../connection-state";
 
 /** The server can be slow to accept when many notebooks are open. */
@@ -40,6 +38,19 @@ async function resolve(url: StreamUrl): Promise<URL | undefined> {
   }
 }
 
+/** Send each request with the runtime's headers, read before each request. */
+export function withHeaders(
+  headers: () => Record<string, string>,
+): typeof globalThis.fetch {
+  return (input, init) => {
+    const sent = new Headers(init?.headers);
+    for (const [name, value] of Object.entries(headers())) {
+      sent.set(name, value);
+    }
+    return globalThis.fetch(input, { ...init, headers: sent });
+  };
+}
+
 /** A resume cursor and replica belong to the session that produced them. */
 function withSessionIdentity(
   url: StreamUrl,
@@ -131,13 +142,7 @@ export function sessionSse(
   return withSessionIdentity(url, (target) =>
     sse(target, {
       parseData: jsonParseWithSpecialChar,
-      fetch: (input, init) => {
-        const sent = new Headers(init?.headers);
-        for (const [name, value] of Object.entries(headers())) {
-          sent.set(name, value);
-        }
-        return globalThis.fetch(input, { ...init, headers: sent });
-      },
+      fetch: withHeaders(headers),
       stop: (frame) =>
         frame.event === "close"
           ? stopOn(frame.data as { code?: number; reason?: string })
@@ -147,7 +152,7 @@ export function sessionSse(
 }
 
 /** A connection that opens and delivers nothing, for static notebooks. */
-const idle: NotebookTransport = ({ signal }, { open }) =>
+export const idle: NotebookTransport = ({ signal }, { open }) =>
   new Promise((settle) => {
     open();
     signal.addEventListener("abort", () => settle(), { once: true });
@@ -183,31 +188,3 @@ export function pyodide(
       signal.addEventListener("abort", () => settle(), { once: true });
     });
 }
-
-/** How the editor follows its notebook. */
-export interface NotebookTransportOptions {
-  static: boolean;
-  transportType: "websocket" | "sse";
-  /** The session route, read before each connection. */
-  url: StreamUrl;
-  /** The runtime's headers, read before each request. */
-  headers: () => Record<string, string>;
-}
-
-/** Choose the transport for the editor's notebook. */
-export function createNotebookTransport(
-  options: NotebookTransportOptions,
-): NotebookTransport {
-  if (options.static) {
-    return idle;
-  }
-  if (isWasm()) {
-    return pyodide((consumer) =>
-      PyodideBridge.INSTANCE.attachMessageConsumer(consumer),
-    );
-  }
-  if (options.transportType === "sse") {
-    return sessionSse(options.url, options.headers);
-  }
-  return websocket(options.url);
-}

package/src/core/websocket/useNotebookClient.ts

@@ -13,12 +13,9 @@ import {
   type NotebookClient,
 } from "../host/notebook-client";
 import { notebookClientAtom } from "./state";
-import {
-  createNotebookTransport,
-  type NotebookTransportOptions,
-} from "./transports/notebook";
+import { createEditorBackend, type EditorBackendOptions } from "./backend";
 
-type Options = Omit<NotebookTransportOptions, "url"> & {
+type Options = Omit<EditorBackendOptions, "url"> & {
   /** The session route, read before each connection. */
   url: () => string;
   /** Wait until the host can accept a connection. */
@@ -33,6 +30,7 @@ export function useNotebookClient(options: Options): NotebookClient {
   const controller = useMemo(() => createNotebookClient(store), [store]);
   const url = useEvent(options.url);
   const headers = useEvent(options.headers);
+  const notebookUrl = useEvent(options.notebookUrl);
   const waitUntilReady = useEvent(options.waitUntilReady);
   const checkHealth = useEvent(options.checkHealth);
   const onNotification = useEvent(options.onNotification);
@@ -42,7 +40,7 @@ export function useNotebookClient(options: Options): NotebookClient {
     let needsInitialReadiness = true;
     store.set(notebookClientAtom, controller.client);
     controller.connect({
-      transport: createNotebookTransport({
+      backend: createEditorBackend({
         static: isStatic,
         transportType,
         url: async () => {
@@ -60,6 +58,7 @@ export function useNotebookClient(options: Options): NotebookClient {
           return url();
         },
         headers,
+        notebookUrl,
       }),
       signal: lifetime.signal,
       checkHealth,
@@ -80,6 +79,7 @@ export function useNotebookClient(options: Options): NotebookClient {
     transportType,
     url,
     headers,
+    notebookUrl,
     waitUntilReady,
     checkHealth,
     onNotification,

package/src/core/websocket/useNotebookSession.tsx

@@ -75,7 +75,6 @@ import {
 import { useStorageActions } from "../storage/state";
 import { useVariablesActions } from "../variables/state";
 import type { VariableName } from "../variables/types";
-import { isWasm } from "../wasm/utils";
 import { WebSocketClosedReason, ConnectionState } from "./types";
 
 const SUPPORTS_LAZY_KERNELS = true;
@@ -418,8 +417,10 @@ export function useNotebookSession(opts: {
         : runtimeManager.getWsURL(sessionId).toString(),
     // SSE uses headers for auth. WebSockets carry the token in the URL.
     headers: () => runtimeManager.headers(),
+    notebookUrl: (notebookId) => runtimeManager.getNotebookURL(notebookId),
+    // Only the host's transports read the URL, so static and WASM never wait.
     waitUntilReady: async () => {
-      if (!isStaticNotebook() && !isWasm() && !runtimeManager.isSameOrigin) {
+      if (!runtimeManager.isSameOrigin) {
         await runtimeManager.waitForHealthy();
       }
     },