@marimo-team/islands 0.25.3-dev137 → 0.25.3-dev139

package/src/core/islands/worker/controller.ts

@@ -1,11 +1,9 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import type { PyodideInterface } from "pyodide";
 import { defaultUserConfig } from "@/core/config/config-schema";
-import type { NotificationPayload } from "@/core/kernel/messages";
 import { DefaultWasmController } from "@/core/wasm/worker/bootstrap";
 import { WasmFileSystem } from "@/core/wasm/worker/fs";
 import type { SerializedBridge } from "@/core/wasm/worker/types";
-import type { JsonString } from "@/utils/json/base64";
 import { Logger } from "../../../utils/Logger";
 
 export class ReadonlyWasmController extends DefaultWasmController {
@@ -40,13 +38,11 @@ export class ReadonlyWasmController extends DefaultWasmController {
   public override async startSession(opts: {
     code: string;
     filename: string | null;
-    onMessage: (message: JsonString<NotificationPayload>) => void;
   }): Promise<SerializedBridge> {
     const bridge = super.startSession({
       queryParameters: {},
       code: opts.code,
       filename: opts.filename,
-      onMessage: opts.onMessage,
       userConfig: defaultUserConfig(),
     });
 

package/src/core/islands/worker/worker.tsx

@@ -8,14 +8,11 @@ import {
   createWorkerParentTransport,
   type RPCSchema,
 } from "rpc-anywhere";
-import type { NotificationPayload } from "@/core/kernel/messages";
 import type { ParentSchema } from "@/core/wasm/rpc";
 import { shouldLoadDuckDBPackages } from "@/core/wasm/utils";
 import { TRANSPORT_ID } from "@/core/wasm/worker/constants";
 import { getPyodideVersion } from "@/core/wasm/worker/getPyodideVersion";
-import { MessageBuffer } from "@/core/wasm/worker/message-buffer";
 import type { RawBridge, SerializedBridge } from "@/core/wasm/worker/types";
-import type { JsonString } from "@/utils/json/base64";
 import { Logger } from "@/utils/Logger";
 import { prettyError } from "../../../utils/errors";
 import { invariant } from "../../../utils/invariant";
@@ -45,14 +42,11 @@ async function loadPyodideAndPackages() {
 }
 
 const pyodideReadyPromise = loadPyodideAndPackages();
-export interface IslandsKernelMessage {
-  message: JsonString<NotificationPayload>;
+/** An event of one session's notebook stream. */
+export interface IslandsNotebookEvent {
+  event: string;
   sessionGeneration: number;
 }
-
-const messageBuffer = new MessageBuffer<IslandsKernelMessage>((payload) => {
-  rpc.send.kernelMessage(payload);
-});
 interface SessionRequest {
   code: string;
   appId: string;
@@ -62,6 +56,8 @@ interface SessionRequest {
 let activeSession:
   | (Omit<SessionRequest, "code"> & { bridge: SerializedBridge })
   | undefined;
+/** Every running session's bridge, kept while the worker follows its notebook. */
+const sessions = new Set<SerializedBridge>();
 let sessionQueue = Promise.resolve();
 
 async function startSession(
@@ -79,18 +75,12 @@ async function startSession(
       code: opts.code,
       filename: `app-${opts.appId}.py`,
     });
-    const nextBridge = await self.controller.startSession({
-      ...notebook,
-      onMessage: (message) => {
-        messageBuffer.push({
-          message,
-          sessionGeneration: opts.sessionGeneration,
-        });
-      },
-    });
-    if (activeSession) {
-      (nextBridge as unknown as PyProxy).destroy();
-    } else {
+    const nextBridge = await self.controller.startSession(notebook);
+    sessions.add(nextBridge);
+    // Follow before the session instantiates. Each subscription begins with
+    // its own kernel-ready and the session's retained notifications.
+    await follow(nextBridge, opts.sessionGeneration);
+    if (!activeSession) {
       activeSession = {
         appId: opts.appId,
         bridge: nextBridge,
@@ -107,10 +97,27 @@ async function startSession(
 }
 
 async function stopActiveSession(): Promise<void> {
-  const session = activeSession;
+  // Stopping closes every session and its notebook subscriptions.
   await self.controller.stopSession();
   activeSession = undefined;
-  (session?.bridge as unknown as PyProxy | undefined)?.destroy();
+  for (const bridge of sessions) {
+    (bridge as unknown as PyProxy).destroy();
+  }
+  sessions.clear();
+}
+
+async function follow(
+  bridge: SerializedBridge,
+  sessionGeneration: number,
+): Promise<void> {
+  const notebook = bridge.notebook;
+  try {
+    await notebook.subscribe(`island-${sessionGeneration}`, null, (event) => {
+      rpc.send.notebookEvent({ event, sessionGeneration });
+    });
+  } finally {
+    notebook.destroy();
+  }
 }
 
 function enqueueSession<T>(operation: () => Promise<T>): Promise<T> {
@@ -243,8 +250,8 @@ export type WorkerSchema = RPCSchema<
     messages: {
       // Emitted when the worker is ready
       ready: {};
-      // Emitted when the kernel sends a message
-      kernelMessage: IslandsKernelMessage;
+      // Emitted for each event of a session's notebook stream
+      notebookEvent: IslandsNotebookEvent;
       // Emitted when the Pyodide is initialized
       initialized: {};
       // Emitted when the Pyodide fails to initialize
@@ -263,13 +270,6 @@ const rpc = createRPC<WorkerSchema, ParentSchema>({
 
 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();
-});
-
 function getMarimoVersion() {
   return self.name; // We store the version in the worker name
 }

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

@@ -51,7 +51,6 @@ vi.mock("@/core/wasm/rpc", () => ({
         readNotebook: mockReadNotebook,
         saveNotebook: mockSaveNotebook,
       },
-      send: { consumerReady: vi.fn() },
     },
     addMessageListener: (event: string, cb: () => void) => {
       rpcListeners[event] = cb;

package/src/core/wasm/bridge.ts

@@ -70,9 +70,6 @@ export class PyodideBridge implements RunRequests, EditRequests {
   private pendingSessionSave: Promise<unknown> = Promise.resolve();
   private notebookConsumers = new Map<string, (event: unknown) => void>();
   private interruptBuffer?: Uint8Array;
-  private messageConsumer:
-    | ((message: MessageEvent<string>) => void)
-    | undefined;
 
   public initialized = new Deferred<void>();
 
@@ -158,9 +155,6 @@ export class PyodideBridge implements RunRequests, EditRequests {
       store.set(wasmInitStateAtom, { kind: "error", message: error });
       this.initialized.reject(new Error(error));
     });
-    this.rpc.addMessageListener("kernelMessage", ({ message }) => {
-      this.messageConsumer?.(new MessageEvent("message", { data: message }));
-    });
   }
 
   private async startSession() {
@@ -212,13 +206,6 @@ export class PyodideBridge implements RunRequests, EditRequests {
     }
   }
 
-  public attachMessageConsumer(
-    consumer: (message: MessageEvent<string>) => void,
-  ) {
-    this.messageConsumer = consumer;
-    this.rpc.proxy.send.consumerReady({});
-  }
-
   public async createNotebookBackend() {
     await this.initialized.promise;
     const { id } = await this.rpc.proxy.request.notebookInfo();

package/src/core/wasm/rpc.ts

@@ -11,7 +11,6 @@ import { TRANSPORT_ID } from "./worker/constants";
 
 export type ParentSchema = RPCSchema<{
   messages: {
-    consumerReady: {};
     bootstrap: WasmRuntimeConfig;
   };
 }>;

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

@@ -2,8 +2,6 @@
 import { loadPyodide, type PyodideInterface } from "pyodide";
 import type { PyCallable, PyProxy } from "pyodide/ffi";
 import type { UserConfig } from "@/core/config/config-schema";
-import type { NotificationPayload } from "@/core/kernel/messages";
-import type { JsonString } from "@/utils/json/base64";
 import { invariant } from "../../../utils/invariant";
 import { Logger } from "../../../utils/Logger";
 import { WasmFileSystem } from "./fs";
@@ -142,20 +140,10 @@ export class DefaultWasmController implements WasmController {
     queryParameters: Record<string, string | string[]>;
     code: string;
     filename: string | null;
-    onMessage: (message: JsonString<NotificationPayload>) => void;
     userConfig: UserConfig;
   }): Promise<SerializedBridge> {
     const sessionGeneration = this.sessionGeneration;
-    const { code, filename, onMessage, queryParameters, userConfig } = opts;
-    // We pass down a messenger object to the code
-    // This is used to have synchronous communication between the JS and Python code
-    // Previously, we used a queue, but this would not properly flush the queue
-    // during processing of a cell's code.
-    //
-    // This adds a messenger object to the global scope (import js; js.messenger.callback)
-    self.messenger = {
-      callback: onMessage,
-    };
+    const { code, filename, queryParameters, userConfig } = opts;
     self.query_params = queryParameters;
     self.user_config = userConfig;
 
@@ -169,14 +157,12 @@ export class DefaultWasmController implements WasmController {
       import js
       from marimo._pyodide.bootstrap import create_session, instantiate
 
-      assert js.messenger, "messenger is not defined"
       assert js.query_params, "query_params is not defined"
 
       def create_session_resources():
         session, bridge = create_session(
           filename="${nbFilename}",
           query_params=js.query_params.to_py(),
-          message_callback=js.messenger.callback,
           user_config=js.user_config.to_py(),
         )
         session_task = None

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

@@ -2,11 +2,7 @@
 import type { PyodideInterface } from "pyodide";
 import type { WasmRuntimeConfig } from "../runtime-config";
 import type { UserConfig } from "@/core/config/config-schema";
-import type {
-  ClientCommand,
-  NotificationPayload,
-} from "@/core/kernel/messages";
-import type { JsonString } from "@/utils/json/base64";
+import type { ClientCommand } from "@/core/kernel/messages";
 import type {
   CodeCompletionRequest,
   CopyNotebookRequest,
@@ -69,7 +65,6 @@ export interface WasmController {
     code: string;
     filename: string | null;
     userConfig: UserConfig;
-    onMessage: (message: JsonString<NotificationPayload>) => void;
   }): Promise<SerializedBridge>;
 }
 

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

@@ -9,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";
@@ -28,7 +25,6 @@ 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,
@@ -73,11 +69,6 @@ async function loadPyodideAndPackages() {
   }
 }
 
-const messageBuffer = new MessageBuffer(
-  (message: JsonString<NotificationPayload>) => {
-    rpc.send.kernelMessage({ message });
-  },
-);
 const bridgeReady = new Deferred<SerializedBridge>();
 let started = false;
 
@@ -155,12 +146,6 @@ const requestHandler = createRPCRequestHandler({
       const startSession = t.wrapAsync(
         self.controller.startSession.bind(self.controller),
       );
-      const initializeOnce = once(() => {
-        rpc.send.initializingMessage({
-          message: "Initializing notebook...",
-        });
-        rpc.send.initialized({});
-      });
       rpc.send.initializingMessage({
         message: "Loading notebook and dependencies...",
       });
@@ -169,12 +154,12 @@ const requestHandler = createRPCRequestHandler({
         filename: opts.filename,
         queryParameters: opts.queryParameters,
         userConfig: opts.userConfig,
-        onMessage: (msg) => {
-          initializeOnce();
-          messageBuffer.push(msg);
-        },
       });
       bridgeReady.resolve(bridge);
+      rpc.send.initializingMessage({
+        message: "Initializing notebook...",
+      });
+      rpc.send.initialized({});
       workerInitSpan.end("ok");
     } catch (error) {
       rpc.send.initializedError({
@@ -385,8 +370,6 @@ export type WorkerSchema = RPCSchema<
       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
@@ -409,13 +392,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

@@ -1,22 +1,23 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { NotebookEvents } from "@marimo-team/client";
+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, consumer: vi.fn() }));
+const wasm = vi.hoisted(() => ({ enabled: false, create: vi.fn() }));
 vi.mock("../../wasm/utils", () => ({ isWasm: () => wasm.enabled }));
 vi.mock("../../wasm/bridge", () => ({
-  PyodideBridge: { INSTANCE: { attachMessageConsumer: wasm.consumer } },
+  PyodideBridge: { INSTANCE: { createNotebookBackend: wasm.create } },
 }));
 
 afterEach(() => {
   wasm.enabled = false;
-  wasm.consumer.mockReset();
+  wasm.create.mockReset();
   vi.unstubAllGlobals();
 });
 
@@ -55,6 +56,8 @@ describe("createEditorBackend", () => {
     ["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);
@@ -62,9 +65,8 @@ describe("createEditorBackend", () => {
     await vi.waitFor(() =>
       expect(
         {
-          idle:
-            events.status === "connected" && !wasm.consumer.mock.calls.length,
-          wasm: wasm.consumer.mock.calls.length === 1,
+          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],
@@ -107,19 +109,45 @@ describe("createEditorBackend", () => {
     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;
+  it("does not edit a static notebook", async () => {
     const fetch = vi.fn();
     vi.stubGlobal("fetch", fetch);
-    const result = await createEditorBackend(chosen)
+    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", () => {

package/src/core/websocket/backend.ts

@@ -9,7 +9,6 @@ import { PyodideBridge } from "../wasm/bridge";
 import { isWasm } from "../wasm/utils";
 import {
   idle,
-  pyodide,
   sessionSse,
   websocket,
   withHeaders,
@@ -61,6 +60,36 @@ function afterStartup({
   };
 }
 
+/**
+ * Follow the notebook running in the WASM worker.
+ *
+ * The worker's backend exists once the worker has started, so the first
+ * connection creates it. Commands are only for a notebook its stream reported.
+ */
+function wasm(): NotebookBackend {
+  let backend: Promise<NotebookBackend> | undefined;
+  let started: NotebookBackend | undefined;
+  return {
+    transport: async (attempt, handlers) => {
+      backend ??= PyodideBridge.INSTANCE.createNotebookBackend();
+      try {
+        started = await backend;
+      } catch (error) {
+        throw new StreamStopped("The WASM notebook did not start", {
+          cause: error,
+        });
+      }
+      return started.transport(attempt, handlers);
+    },
+    notebook: (id) => {
+      if (started === undefined) {
+        throw new Error("The WASM notebook has not started");
+      }
+      return started.notebook(id);
+    },
+  };
+}
+
 /** Commands for a notebook the editor cannot change through its backend. */
 function uneditable(): ReturnType<NotebookBackend["notebook"]> {
   return {
@@ -84,12 +113,7 @@ export function createEditorBackend(
     return { transport: idle, notebook: uneditable };
   }
   if (isWasm()) {
-    return {
-      transport: pyodide((consumer) =>
-        PyodideBridge.INSTANCE.attachMessageConsumer(consumer),
-      ),
-      notebook: uneditable,
-    };
+    return wasm();
   }
   const url = afterStartup(options);
   return {

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

@@ -11,7 +11,7 @@ import { createNotebookReplica } from "../../../host/replica";
 import { fakeBackend, fixtureMessages } from "../../../host/__tests__/fixtures";
 import { ConnectionState } from "../../../host/connection";
 import { Deferred } from "@/utils/Deferred";
-import { pyodide, sessionSse, websocket } from "../notebook";
+import { sessionSse, websocket } from "../notebook";
 
 const encoder = new TextEncoder();
 
@@ -403,43 +403,3 @@ describe.each(["websocket", "sse"] as const)(
     });
   },
 );
-
-describe("pyodide", () => {
-  it("delivers worker notifications as marimo events and host events as they are", async () => {
-    let consumer: ((message: MessageEvent<string>) => void) | undefined;
-    const events = new NotebookEvents(
-      pyodide((next) => {
-        consumer = next;
-      }),
-      { types: ["notebook", "marimo"] },
-    );
-    const received = collect(events);
-    await vi.waitFor(() => expect(consumer).toBeDefined());
-    const send = (data: string) =>
-      consumer?.(new MessageEvent("message", { data }));
-    send('{"data":{"op":"kernel-ready"}}');
-    send('{"event":"marimo","id":"3","data":{"marimo":{"op":"reconnected"}}}');
-    expect(received).toMatchInlineSnapshot(`
-      [
-        {
-          "data": {
-            "marimo": {
-              "op": "kernel-ready",
-            },
-          },
-          "event": "marimo",
-        },
-        {
-          "data": {
-            "marimo": {
-              "op": "reconnected",
-            },
-          },
-          "event": "marimo",
-          "id": "3",
-        },
-      ]
-    `);
-    events.close();
-  });
-});

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

@@ -157,34 +157,3 @@ export const idle: NotebookTransport = ({ signal }, { open }) =>
     open();
     signal.addEventListener("abort", () => settle(), { once: true });
   });
-
-/**
- * Follow the notebook running in the WASM worker.
- *
- * The worker sends marimo notifications as `{ data: notification }`, which
- * arrive as `marimo` events.
- */
-export function pyodide(
-  attach: (consumer: (message: MessageEvent<string>) => void) => void,
-): NotebookTransport {
-  return ({ signal }, { open, message }) =>
-    new Promise((settle) => {
-      open();
-      attach(({ data }) => {
-        if (signal.aborted) {
-          return;
-        }
-        const value = jsonParseWithSpecialChar(data);
-        // The worker sends notifications without the host's event envelope.
-        message(
-          typeof value === "object" &&
-            value !== null &&
-            !("event" in value) &&
-            "data" in value
-            ? { event: "marimo", data: { marimo: value.data } }
-            : value,
-        );
-      });
-      signal.addEventListener("abort", () => settle(), { once: true });
-    });
-}

package/dist/assets/__vite-browser-external-SaDYcg7x.js

@@ -1 +0,0 @@
-import{t as e}from"./worker-CfME3uyi.js";var t=e(((e,t)=>{t.exports={}}));export default t();