@marimo-team/islands 0.25.3-dev14 → 0.25.3-dev140

package/src/core/wasm/worker/worker.ts

@@ -1,5 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { unstable_NotebookCall } from "@marimo-team/client";
 import type { PyodideInterface } from "pyodide";
 import {
   createRPC,
@@ -8,14 +9,11 @@ import {
   type RPCSchema,
 } from "rpc-anywhere";
 import type { UserConfig } from "@/core/config/config-schema";
-import type { NotificationPayload } from "@/core/kernel/messages";
 import type {
   ListPackagesResponse,
   PackageOperationResponse,
   SaveNotebookRequest,
 } from "@/core/network/types";
-import type { JsonString } from "@/utils/json/base64";
-import { once } from "@/utils/once";
 import { decodeUtf8 } from "@/utils/strings";
 import { Deferred } from "../../../utils/Deferred";
 import { prettyError } from "../../../utils/errors";
@@ -27,12 +25,11 @@ import { TRANSPORT_ID } from "./constants";
 import { WasmFileSystem } from "./fs";
 import { getController } from "./getController";
 import { getPyodideVersion } from "./getPyodideVersion";
-import { MessageBuffer } from "./message-buffer";
 import { t } from "./tracer";
 import type {
   BridgePayload,
   RawBridge,
-  SerializedBridge,
+  WasmSession,
   WasmController,
 } from "./types";
 import { shouldLoadDuckDBPackages } from "../utils";
@@ -72,16 +69,74 @@ async function loadPyodideAndPackages() {
   }
 }
 
-const messageBuffer = new MessageBuffer(
-  (message: JsonString<NotificationPayload>) => {
-    rpc.send.kernelMessage({ message });
-  },
-);
-const bridgeReady = new Deferred<SerializedBridge>();
+const sessionReady = new Deferred<WasmSession>();
 let started = false;
 
 // Handle RPC requests
+const notebookSubscriptions = new Map<
+  string,
+  { promise: Promise<void>; cancel(): void }
+>();
+async function getNotebook() {
+  const session = await sessionReady.promise;
+  return session.bridge.notebook;
+}
+
 const requestHandler = createRPCRequestHandler({
+  notebookInfo: async () => {
+    const notebook = await getNotebook();
+    try {
+      return { id: notebook.id };
+    } finally {
+      notebook.destroy();
+    }
+  },
+  notebookSubscribe: (request: { id: string; cursor: string | null }) => {
+    let canceled = false;
+    const subscribing = (async () => {
+      const session = await sessionReady.promise;
+      const notebook = session.bridge.notebook;
+      try {
+        await notebook.subscribe(request.id, request.cursor, (event) => {
+          rpc.send.notebookEvent({ id: request.id, event });
+        });
+        if (!canceled) {
+          void session.start().catch((error: unknown) => {
+            rpc.send.initializedError({ error: prettyError(error) });
+          });
+        }
+      } finally {
+        notebook.destroy();
+      }
+    })();
+    notebookSubscriptions.set(request.id, {
+      promise: subscribing,
+      cancel: () => {
+        canceled = true;
+      },
+    });
+    return subscribing.finally(() => notebookSubscriptions.delete(request.id));
+  },
+  notebookUnsubscribe: async (id: string) => {
+    // An RPC timeout can close a reader while its subscription is still opening.
+    const pending = notebookSubscriptions.get(id);
+    pending?.cancel();
+    await pending?.promise.catch(() => undefined);
+    const notebook = await getNotebook();
+    try {
+      notebook.unsubscribe(id);
+    } finally {
+      notebook.destroy();
+    }
+  },
+  notebookCall: async (call: unstable_NotebookCall): Promise<unknown> => {
+    const notebook = await getNotebook();
+    try {
+      return JSON.parse(await notebook.call(JSON.stringify(call)));
+    } finally {
+      notebook.destroy();
+    }
+  },
   /**
    * Start the session
    */
@@ -105,29 +160,23 @@ const requestHandler = createRPCRequestHandler({
         code: opts.code,
         filename: opts.filename,
       });
-      const startSession = t.wrapAsync(
-        self.controller.startSession.bind(self.controller),
+      const prepareSession = t.wrapAsync(
+        self.controller.prepareSession.bind(self.controller),
       );
-      const initializeOnce = once(() => {
-        rpc.send.initializingMessage({
-          message: "Initializing notebook...",
-        });
-        rpc.send.initialized({});
-      });
       rpc.send.initializingMessage({
         message: "Loading notebook and dependencies...",
       });
-      const bridge = await startSession({
+      const session = await prepareSession({
         code: opts.code,
         filename: opts.filename,
         queryParameters: opts.queryParameters,
         userConfig: opts.userConfig,
-        onMessage: (msg) => {
-          initializeOnce();
-          messageBuffer.push(msg);
-        },
       });
-      bridgeReady.resolve(bridge);
+      sessionReady.resolve(session);
+      rpc.send.initializingMessage({
+        message: "Initializing notebook...",
+      });
+      rpc.send.initialized({});
       workerInitSpan.end("ok");
     } catch (error) {
       rpc.send.initializedError({
@@ -254,7 +303,7 @@ const requestHandler = createRPCRequestHandler({
    */
   saveNotebook: async (opts: SaveNotebookRequest) => {
     await pyodideReadyPromise; // Make sure loading is done
-    const bridge = await bridgeReady.promise;
+    const { bridge } = await sessionReady.promise;
     await bridge.save(JSON.stringify(opts));
   },
 
@@ -302,7 +351,7 @@ const requestHandler = createRPCRequestHandler({
     }
 
     // Perform the function call to the Python bridge
-    const bridge = await bridgeReady.promise;
+    const { bridge } = await sessionReady.promise;
 
     // Serialize the payload
     const payloadString =
@@ -335,10 +384,9 @@ const requestHandler = createRPCRequestHandler({
 export type WorkerSchema = RPCSchema<
   {
     messages: {
+      notebookEvent: { id: string; event: string };
       // Emitted when the worker is ready
       ready: {};
-      // Emitted when the kernel sends a message
-      kernelMessage: { message: JsonString<NotificationPayload> };
       // Emitted when the Pyodide is initialized
       initialized: {};
       // Emitted when the Pyodide is initializing, with new messages
@@ -361,13 +409,6 @@ rpc.addMessageListener("bootstrap", (config) => runtimeConfig.resolve(config));
 
 rpc.send("ready", {});
 
-/// Listeners
-// When the consumer is ready, start the message buffer
-rpc.addMessageListener("consumerReady", async () => {
-  await pyodideReadyPromise; // Make sure loading is done
-  messageBuffer.start();
-});
-
 const namesThatRequireSync = new Set<keyof RawBridge>([
   "save",
   "save_app_config",

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

@@ -0,0 +1,228 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { NotebookEvents, StreamStopped } from "@marimo-team/client";
+import { createStore } from "jotai/vanilla";
+import { afterEach, describe, expect, it, onTestFinished, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import { ConnectionState } from "../../host/connection";
+import { type FakeEdit, fakeBackend } from "../../host/__tests__/fixtures";
+import { createNotebookReplica } from "../../host/replica";
+import { createEditorBackend } from "../backend";
+
+const wasm = vi.hoisted(() => ({ enabled: false, create: vi.fn() }));
+vi.mock("../../wasm/utils", () => ({ isWasm: () => wasm.enabled }));
+vi.mock("../../wasm/bridge", () => ({
+  PyodideBridge: { INSTANCE: { createNotebookBackend: wasm.create } },
+}));
+
+afterEach(() => {
+  wasm.enabled = false;
+  wasm.create.mockReset();
+  vi.unstubAllGlobals();
+});
+
+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}`),
+  waitUntilReady: async () => undefined,
+  checkHealth: async () => true,
+};
+
+/** Record each socket the transport opens. They never connect. */
+function sockets() {
+  const opened: string[] = [];
+  vi.stubGlobal(
+    "WebSocket",
+    class extends EventTarget {
+      public constructor(url: URL) {
+        super();
+        opened.push(String(url));
+      }
+      public close() {}
+    },
+  );
+  return opened;
+}
+
+describe("createEditorBackend", () => {
+  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 worker = vi.fn(() => new Promise<void>(() => undefined));
+    wasm.create.mockResolvedValue(fakeBackend({ transport: worker }));
+    const fetch = vi.fn(() => new Promise<Response>(() => undefined));
+    const opened = sockets();
+    vi.stubGlobal("fetch", fetch);
+    const events = new NotebookEvents(createEditorBackend(chosen).transport);
+    await vi.waitFor(() =>
+      expect(
+        {
+          idle: events.status === "connected" && !wasm.create.mock.calls.length,
+          wasm: worker.mock.calls.length === 1,
+          fetch: fetch.mock.calls.length === 1,
+          socket: opened.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("does not edit a static notebook", async () => {
+    const fetch = vi.fn();
+    vi.stubGlobal("fetch", fetch);
+    const result = await createEditorBackend({ ...options, static: true })
+      .notebook("n1")
+      .edit([], { idempotencyKey: "edit-1" });
+    expect(result).toMatchObject({ ok: false, status: 501 });
+    expect(fetch).not.toHaveBeenCalled();
+  });
+
+  it("edits the WASM notebook through the worker once its stream starts", async () => {
+    wasm.enabled = true;
+    const edit = vi.fn<FakeEdit>();
+    const transport = vi.fn(async () => undefined);
+    wasm.create.mockResolvedValue(fakeBackend({ transport, edit }));
+    const backend = createEditorBackend(options);
+    expect(wasm.create).not.toHaveBeenCalled();
+    expect(() => backend.notebook("n1")).toThrow("has not started");
+
+    const attempt = { cursor: null, signal: new AbortController().signal };
+    const handlers = { open: vi.fn(), message: vi.fn() };
+    await backend.transport(attempt, handlers);
+    await backend.transport(attempt, handlers);
+    expect(wasm.create).toHaveBeenCalledOnce();
+    expect(transport).toHaveBeenCalledWith(attempt, handlers);
+    await backend.notebook("n1").edit([], { idempotencyKey: "edit-1" });
+    expect(edit).toHaveBeenCalledWith("n1", [], { idempotencyKey: "edit-1" });
+  });
+
+  it("stops following WASM when the worker does not start", async () => {
+    wasm.enabled = true;
+    wasm.create.mockRejectedValue(new Error("Initialization failed"));
+    await expect(
+      createEditorBackend(options).transport(
+        { cursor: null, signal: new AbortController().signal },
+        { open: vi.fn(), message: vi.fn() },
+      ),
+    ).rejects.toBeInstanceOf(StreamStopped);
+  });
+});
+
+describe("the session backend's startup", () => {
+  it.each([true, false])(
+    "rechecks health after failed startup (healthy: %s)",
+    async (healthy) => {
+      const opened = sockets();
+      const waitUntilReady = vi
+        .fn()
+        .mockRejectedValue(new Error("Startup gave up"));
+      const checkHealth = vi.fn().mockResolvedValue(healthy);
+      const store = createStore();
+      const client = createNotebookReplica(store);
+      onTestFinished(
+        client.connect(
+          createEditorBackend({ ...options, waitUntilReady, checkHealth }),
+          { onNotification: vi.fn() },
+        ),
+      );
+      await vi.waitFor(() =>
+        expect(store.get(client.connectionAtom).state).toBe(
+          ConnectionState.CLOSED,
+        ),
+      );
+
+      await client.reconnect();
+      await vi.waitFor(() => expect(opened).toHaveLength(healthy ? 1 : 0));
+      expect({
+        healthChecks: checkHealth.mock.calls.length,
+        startupWaits: waitUntilReady.mock.calls.length,
+        state: store.get(client.connectionAtom).state,
+      }).toEqual({
+        healthChecks: 1,
+        startupWaits: 1,
+        state: healthy ? ConnectionState.CONNECTING : ConnectionState.CLOSED,
+      });
+    },
+  );
+
+  it("shares one startup wait across connections", async () => {
+    const opened = sockets();
+    const ready = new Deferred<void>();
+    const waitUntilReady = vi.fn(() => ready.promise);
+    const backend = createEditorBackend({ ...options, waitUntilReady });
+    const client = createNotebookReplica(createStore());
+    // Strict Mode connects, disconnects, and connects again.
+    client.connect(backend, { onNotification: vi.fn() })();
+    onTestFinished(client.connect(backend, { onNotification: vi.fn() }));
+    await Promise.resolve();
+    expect(opened).toEqual([]);
+
+    ready.resolve();
+    await vi.waitFor(() => expect(opened).toHaveLength(1));
+    expect(waitUntilReady).toHaveBeenCalledOnce();
+  });
+
+  it("opens nothing for a connection that ended during startup", async () => {
+    const opened = sockets();
+    const ready = new Deferred<void>();
+    const store = createStore();
+    const client = createNotebookReplica(store);
+    const disconnect = client.connect(
+      createEditorBackend({ ...options, waitUntilReady: () => ready.promise }),
+      { onNotification: vi.fn() },
+    );
+    expect(store.get(client.connectionAtom)).toEqual({
+      state: ConnectionState.CONNECTING,
+    });
+    disconnect();
+    ready.resolve();
+    await ready.promise;
+    await Promise.resolve();
+    expect(opened).toEqual([]);
+    expect(store.get(client.connectionAtom)).toEqual({
+      state: ConnectionState.NOT_STARTED,
+    });
+  });
+});

package/src/core/websocket/__tests__/connection-state.test.ts

@@ -64,16 +64,4 @@ describe("classifyCloseEvent", () => {
       });
     });
   });
-
-  describe("transport exhaustion", () => {
-    it("MARIMO_TRANSPORT_EXHAUSTED → gave-up with KERNEL_DISCONNECTED", () => {
-      const decision = classify("MARIMO_TRANSPORT_EXHAUSTED");
-      expect(decision.kind).toBe("gave-up");
-      expect(decision.status).toEqual({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-        reason: "kernel not found",
-      });
-    });
-  });
 });

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

@@ -0,0 +1,141 @@
+/* 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 type { NotebookTransport } from "@marimo-team/client";
+import { ConnectionState } from "../types";
+import { createEditorBackend } from "../backend";
+import { useNotebookReplica } from "../useNotebookReplica";
+import { fakeBackend, fixtureMessages } from "../../host/__tests__/fixtures";
+import { agentAttachmentsAtom } from "../../attachments/state";
+import { notebookReplicaAtom } from "../state";
+
+vi.mock("../backend", () => ({ createEditorBackend: vi.fn() }));
+
+/** A host-framed native notification. */
+const reconnected = {
+  event: "marimo",
+  data: { marimo: { op: "reconnected" } },
+};
+
+/** A transport whose connections the test drives. */
+function makeTransport() {
+  const opened: {
+    signal: AbortSignal;
+    open(): void;
+    receive(message: unknown): void;
+  }[] = [];
+  const transport: NotebookTransport = ({ signal }, handlers) =>
+    new Promise((resolve) => {
+      opened.push({ signal, open: handlers.open, receive: handlers.message });
+      signal.addEventListener("abort", () => resolve(), { once: true });
+    });
+  return {
+    transport,
+    opened,
+    /** Open the latest connection and deliver a host-framed notification. */
+    notify: () => opened.at(-1)?.receive(reconnected),
+  };
+}
+
+function setup(strict = false) {
+  const store = createStore();
+  const options = {
+    static: false,
+    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(),
+  };
+  const wrapper = ({ children }: PropsWithChildren) => (
+    <Provider store={store}>{children}</Provider>
+  );
+  const hook = renderHook((props) => useNotebookReplica(props), {
+    initialProps: options,
+    reactStrictMode: strict,
+    wrapper,
+  });
+  return { ...hook, options, store };
+}
+
+beforeEach(() => {
+  vi.mocked(createEditorBackend).mockReset();
+});
+
+describe("useNotebookReplica", () => {
+  it("exposes host presence to the UI without waiting for kernel-ready", async () => {
+    const fake = makeTransport();
+    vi.mocked(createEditorBackend).mockReturnValue(
+      fakeBackend({ transport: fake.transport }),
+    );
+    const { result, store, unmount } = setup();
+    await act(async () => Promise.resolve());
+    expect(store.get(notebookReplicaAtom)).toBe(result.current);
+    const [connection] = fake.opened;
+    act(() => {
+      connection.open();
+      fixtureMessages("notebook-events.sse").forEach(connection.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(notebookReplicaAtom)).toBeNull();
+    expect(store.get(agentAttachmentsAtom)).toEqual([]);
+  });
+
+  it("uses current callbacks and URL providers without reconnecting on a rerender", async () => {
+    const fake = makeTransport();
+    vi.mocked(createEditorBackend).mockReturnValue(
+      fakeBackend({ transport: fake.transport }),
+    );
+    const { rerender, options, unmount } = setup();
+    await act(async () => Promise.resolve());
+    const onNotification = vi.fn();
+    rerender({ ...options, onNotification, url: () => "ws://redirected/ws" });
+    act(() => fake.notify());
+    expect(options.onNotification).not.toHaveBeenCalled();
+    expect(onNotification).toHaveBeenCalledExactlyOnceWith({
+      data: { op: "reconnected" },
+    });
+    const [{ url }] = vi.mocked(createEditorBackend).mock.calls[0];
+    expect(url()).toBe("ws://redirected/ws");
+    expect(fake.opened).toHaveLength(1);
+    unmount();
+  });
+
+  it("follows only the active connection under Strict Mode", async () => {
+    const fake = makeTransport();
+    vi.mocked(createEditorBackend).mockReturnValue(
+      fakeBackend({ transport: fake.transport }),
+    );
+    const { result, store, options, unmount } = setup(true);
+    await act(async () => Promise.resolve());
+    const live = fake.opened.filter((attempt) => !attempt.signal.aborted);
+    expect(live).toHaveLength(1);
+    const [active] = live;
+    act(() => active.open());
+    expect(store.get(result.current.connectionAtom)).toEqual({
+      state: ConnectionState.OPEN,
+    });
+    act(() => fake.opened.forEach((attempt) => attempt.receive(reconnected)));
+    expect(options.onNotification).toHaveBeenCalledOnce();
+    unmount();
+    expect(active.signal.aborted).toBe(true);
+    const attempts = fake.opened.length;
+    await result.current.reconnect();
+    expect(fake.opened).toHaveLength(attempts);
+  });
+});