@marimo-team/islands 0.25.3-dev141 → 0.25.3-dev143

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

@@ -0,0 +1,187 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { PyodideInterface } from "pyodide";
+import { loadPyodide } from "pyodide";
+import { afterEach, expect, it, vi } from "vitest";
+import { defaultUserConfig } from "@/core/config/config-schema";
+import { Deferred } from "@/utils/Deferred";
+import { DefaultWasmController } from "../worker/bootstrap";
+import { getController } from "../worker/getController";
+import type { WasmRuntimeConfig } from "../runtime-config";
+
+interface WorkerRequests {
+  startSession(opts: {
+    queryParameters: Record<string, string | string[]>;
+    code: string;
+    filename: string;
+    userConfig: ReturnType<typeof defaultUserConfig>;
+  }): Promise<void>;
+  notebookSubscribe(opts: { id: string; cursor: string | null }): Promise<void>;
+  notebookUnsubscribe(id: string): Promise<void>;
+}
+
+const { rpc, handlers } = vi.hoisted(() => ({
+  handlers: {} as {
+    requests: WorkerRequests;
+    bootstrap: (config: WasmRuntimeConfig) => void;
+  },
+  rpc: {
+    send: Object.assign(vi.fn(), {
+      initialized: vi.fn(),
+      initializedError: vi.fn(),
+      initializingMessage: vi.fn(),
+      notebookEvent: vi.fn(),
+    }),
+  },
+}));
+
+vi.mock("rpc-anywhere", () => ({
+  createRPCRequestHandler: (requests: WorkerRequests) => {
+    handlers.requests = requests;
+    return requests;
+  },
+  createWorkerParentTransport: vi.fn(),
+  createRPC: () => ({
+    ...rpc,
+    addMessageListener: (
+      _name: string,
+      listener: (config: WasmRuntimeConfig) => void,
+    ) => {
+      handlers.bootstrap = listener;
+    },
+  }),
+}));
+vi.mock("pyodide", () => ({ loadPyodide: vi.fn(), version: "314.0.0" }));
+vi.mock("../worker/getController", () => ({ getController: vi.fn() }));
+vi.mock("../worker/getMarimoWheel", () => ({
+  getMarimoWheel: () => "marimo-base",
+}));
+
+afterEach(() => {
+  vi.restoreAllMocks();
+  vi.clearAllMocks();
+  vi.resetModules();
+});
+
+async function startWorker(subscribing = Promise.resolve()) {
+  const subscribers = new Map<string, (event: string) => void>();
+  const event = JSON.stringify({
+    event: "marimo",
+    data: { marimo: { op: "query-params-set", key: "startup", value: "ran" } },
+  });
+  const notebook = {
+    subscribe: vi.fn(
+      async (id: string, _cursor: unknown, emit: (event: string) => void) => {
+        await subscribing;
+        subscribers.set(id, emit);
+      },
+    ),
+    unsubscribe: (id: string) => subscribers.delete(id),
+    set_persistence: vi.fn(),
+    destroy: vi.fn(),
+  };
+  const bridge = { notebook, destroy: vi.fn() };
+  const init = Object.assign(
+    vi.fn(() => subscribers.forEach((emit) => emit(event))),
+    { destroy: vi.fn() },
+  );
+  const stop = Object.assign(
+    vi.fn(async () => subscribers.clear()),
+    {
+      destroy: vi.fn(),
+    },
+  );
+  const packages = { toJs: () => [], destroy: vi.fn() };
+  const resources = Object.assign([bridge, init, packages, stop], {
+    destroy: vi.fn(),
+  });
+  const pyodide = {
+    runPython: vi.fn(() => resources),
+    loadPackagesFromImports: vi.fn(async () => undefined),
+    loadedPackages: {},
+  };
+  vi.mocked(loadPyodide).mockResolvedValue(
+    pyodide as unknown as PyodideInterface,
+  );
+  const controller = new DefaultWasmController();
+  vi.spyOn(controller, "mountFilesystem").mockResolvedValue({
+    code: "",
+    filename: "notebook.py",
+  });
+  vi.mocked(getController).mockResolvedValue(controller);
+  await import("../worker/worker");
+  handlers.bootstrap({ version: "0.25.1" });
+  await handlers.requests.startSession({
+    code: "",
+    filename: "notebook.py",
+    queryParameters: {},
+    userConfig: defaultUserConfig(),
+  });
+  return { controller, init, notebook, event, pyodide };
+}
+
+it("delivers startup notifications after a delayed first subscription, without rerunning on resubscribe", async () => {
+  const { controller, init, notebook, event, pyodide } = await startWorker();
+  try {
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(rpc.send.initialized).toHaveBeenCalledOnce();
+    expect(notebook.set_persistence).toHaveBeenCalledExactlyOnceWith(
+      expect.any(Function),
+    );
+    expect(pyodide.loadPackagesFromImports).toHaveResolved();
+    expect(init).not.toHaveBeenCalled();
+
+    await handlers.requests.notebookSubscribe({ id: "first", cursor: null });
+    await vi.waitFor(() => expect(init).toHaveBeenCalledOnce());
+    expect(rpc.send.notebookEvent.mock.calls).toEqual([
+      [{ id: "first", event }],
+    ]);
+
+    await handlers.requests.notebookUnsubscribe("first");
+    await handlers.requests.notebookSubscribe({ id: "second", cursor: null });
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(init).toHaveBeenCalledOnce();
+    expect(notebook.set_persistence).toHaveBeenCalledOnce();
+  } finally {
+    await controller.stopSession();
+  }
+});
+
+it("does not authorize startup through a subscription canceled while opening", async () => {
+  const subscribing = new Deferred<void>();
+  const { controller, init, notebook } = await startWorker(subscribing.promise);
+  try {
+    const opened = handlers.requests.notebookSubscribe({
+      id: "canceled",
+      cursor: null,
+    });
+    await vi.waitFor(() => expect(notebook.subscribe).toHaveBeenCalledOnce());
+    const closed = handlers.requests.notebookUnsubscribe("canceled");
+    subscribing.resolve();
+    await Promise.all([opened, closed]);
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(init).not.toHaveBeenCalled();
+
+    await handlers.requests.notebookSubscribe({ id: "active", cursor: null });
+    await vi.waitFor(() => expect(init).toHaveBeenCalledOnce());
+  } finally {
+    await controller.stopSession();
+  }
+});
+
+it("waits for a successful subscription after an opening failure", async () => {
+  const { controller, init, notebook } = await startWorker();
+  try {
+    notebook.subscribe.mockRejectedValueOnce(new Error("Subscription failed"));
+    await expect(
+      handlers.requests.notebookSubscribe({ id: "failed", cursor: null }),
+    ).rejects.toThrow("Subscription failed");
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(init).not.toHaveBeenCalled();
+
+    await handlers.requests.notebookSubscribe({ id: "active", cursor: null });
+    await vi.waitFor(() => expect(init).toHaveBeenCalledOnce());
+  } finally {
+    await controller.stopSession();
+  }
+});

package/src/core/wasm/bridge.ts

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 /* oxlint-disable typescript/no-explicit-any */
 
+import { createRPCRequestHandler } from "rpc-anywhere";
 import { toast } from "@/components/ui/use-toast";
 import { userConfigAtom } from "@/core/config/config";
 import { serializeBlob } from "@/utils/blob";
@@ -8,9 +9,12 @@ import { Deferred } from "@/utils/Deferred";
 import { throwNotImplemented } from "@/utils/functions";
 import { Logger } from "@/utils/Logger";
 import { reloadSafe } from "@/utils/reload-safe";
+import { once } from "@/utils/once";
 import { generateUUID } from "@/utils/uuid";
 import { createModuleWorker } from "@/utils/worker";
 import { notebookIsRunningAtom } from "../cells/cells";
+import { flushDocumentChanges } from "../cells/document-changes";
+import { createWasmNotebookBackend } from "../host";
 import type { ClientCommand } from "../kernel/messages";
 import { getInitialAppMode } from "../mode";
 import { API } from "../network/api";
@@ -38,10 +42,9 @@ import type {
 } from "../network/types";
 import { filenameAtom } from "../saving/file-state";
 import { store } from "../state/jotai";
-import { BasicTransport } from "../websocket/transports/basic";
-import type { IConnectionTransport } from "../websocket/transports/transport";
+import { notebookReplicaAtom } from "../websocket/state";
 import { PyodideRouter } from "./router";
-import { getWorkerRPC } from "./rpc";
+import { getWorkerRPC, type PersistNotebookRequest } from "./rpc";
 import { getWasmRuntimeConfig } from "./runtime-config";
 import { createShareableLink } from "./share";
 import { wasmInitStateAtom } from "./state";
@@ -68,11 +71,8 @@ export class PyodideBridge implements RunRequests, EditRequests {
 
   private rpc!: ReturnType<typeof getWorkerRPC<WorkerSchema>>;
   private saveRpc: SaveWorker | undefined;
-  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>();
 
@@ -82,7 +82,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
       return {
         readFile: throwNotImplemented,
         readNotebook: async () => (await notebookFileStore.readFile()) ?? "",
-        saveNotebook: throwNotImplemented,
+        writeNotebook: throwNotImplemented,
       };
     }
 
@@ -97,7 +97,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
       },
     );
 
-    const rpc = getWorkerRPC<SaveWorkerSchema>(saveWorker);
+    const rpc = getWorkerRPC<SaveWorkerSchema>(saveWorker, undefined, Infinity);
     rpc.send.bootstrap(runtimeConfig);
     return rpc.proxy.request;
   }
@@ -116,9 +116,28 @@ export class PyodideBridge implements RunRequests, EditRequests {
     });
 
     // Create the RPC
-    this.rpc = getWorkerRPC<WorkerSchema>(worker);
+    this.rpc = getWorkerRPC<WorkerSchema>(
+      worker,
+      createRPCRequestHandler({
+        persistNotebook: async (snapshot: PersistNotebookRequest) => {
+          if (!this.saveRpc) {
+            throw new Error("Notebook storage is not initialized");
+          }
+          await this.saveRpc.writeNotebook(snapshot);
+          notebookFileStore.saveFile(snapshot.contents);
+          fallbackFileStore.saveFile(snapshot.contents);
+          return null;
+        },
+      }),
+      // Synchronous Python execution can block worker requests indefinitely.
+      Infinity,
+    );
     this.rpc.send.bootstrap(runtimeConfig);
 
+    this.rpc.addMessageListener("notebookEvent", ({ id, event }) => {
+      this.notebookConsumers.get(id)?.(JSON.parse(event));
+    });
+
     // Listeners
     this.rpc.addMessageListener("ready", () => {
       this.startSession();
@@ -154,9 +173,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() {
@@ -208,11 +224,31 @@ 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;
+    return createWasmNotebookBackend({
+      subscribe: async (cursor, emit, signal) => {
+        signal.throwIfAborted();
+        const id = generateUUID();
+        this.notebookConsumers.set(id, emit);
+        const close = once(() => {
+          signal.removeEventListener("abort", close);
+          this.notebookConsumers.delete(id);
+          void this.rpc.proxy.request
+            .notebookUnsubscribe(id)
+            .catch(Logger.error);
+        });
+        signal.addEventListener("abort", close, { once: true });
+        try {
+          await this.rpc.proxy.request.notebookSubscribe({ id, cursor });
+          return close;
+        } catch (error) {
+          close();
+          throw error;
+        }
+      },
+      call: (call) => this.rpc.proxy.request.notebookCall(call),
+    });
   }
 
   public sendRename: EditRequests["sendRename"] = async ({ filename }) => {
@@ -230,30 +266,23 @@ export class PyodideBridge implements RunRequests, EditRequests {
     return null;
   };
 
-  public sendSave: EditRequests["sendSave"] = async (request) => {
-    if (!this.saveRpc) {
-      Logger.warn("Save RPC not initialized");
-      return null;
+  private async saveHostNotebook() {
+    const replica = store.get(notebookReplicaAtom);
+    const notebook = replica && store.get(replica.notebookAtom);
+    const backend = replica && store.get(replica.backendAtom);
+    if (!notebook || !backend) {
+      throw new Error("The notebook is not connected");
     }
-
-    const durableSave = this.saveRpc.saveNotebook(request);
-    // Generated exports read the session-owned app in the main worker.
-    this.pendingSessionSave = this.pendingSessionSave
-      .catch(() => undefined)
-      .then(async () => {
-        await durableSave;
-        await this.rpc.proxy.request.saveNotebook(request);
-      });
-    void this.pendingSessionSave.catch((error) => {
-      Logger.error(error);
-    });
-
-    await durableSave;
-    const code = await this.readCode();
-    if (code.contents) {
-      notebookFileStore.saveFile(code.contents);
-      fallbackFileStore.saveFile(code.contents);
+    const result = await backend.notebook(notebook.id).save();
+    if (!result.ok) {
+      throw new Error(result.problem.detail ?? result.problem.title);
     }
+  }
+
+  public sendSave: EditRequests["sendSave"] = async ({ layout }) => {
+    await flushDocumentChanges();
+    await this.rpc.proxy.request.setNotebookLayout(layout ?? null);
+    await this.saveHostNotebook();
     return null;
   };
 
@@ -370,6 +399,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
       functionName: "save_app_config",
       payload: request,
     });
+    await this.saveHostNotebook();
     return null;
   };
 
@@ -409,7 +439,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
     return response as Snippets;
   };
 
-  public openFile: EditRequests["openFile"] = async ({ path }) => {
+  public openFile: EditRequests["openFile"] = async () => {
     const url = createShareableLink({
       code: null,
       baseUrl: window.location.origin,
@@ -453,7 +483,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
     return null;
   };
 
-  public sendInstantiate: RunRequests["sendInstantiate"] = async (request) => {
+  public sendInstantiate: RunRequests["sendInstantiate"] = async () => {
     return null;
   };
 
@@ -535,7 +565,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
   public exportAsHTML: EditRequests["exportAsHTML"] = async (
     request: ExportAsHTMLRequest,
   ) => {
-    await this.pendingSessionSave;
+    await flushDocumentChanges();
     const response = await this.rpc.proxy.request.bridge({
       functionName: "export_html",
       payload: withDevAssetUrl(request),
@@ -546,7 +576,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
   public exportAsMarkdown: EditRequests["exportAsMarkdown"] = async (
     request: ExportAsMarkdownRequest,
   ) => {
-    await this.pendingSessionSave;
+    await flushDocumentChanges();
     const response = await this.rpc.proxy.request.bridge({
       functionName: "export_markdown",
       payload: request,
@@ -557,7 +587,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
   public exportAsScript: EditRequests["exportAsScript"] = async (
     request: ExportAsScriptRequest,
   ) => {
-    await this.pendingSessionSave;
+    await flushDocumentChanges();
     const response = await this.rpc.proxy.request.bridge({
       functionName: "export_script",
       payload: request,
@@ -628,7 +658,12 @@ export class PyodideBridge implements RunRequests, EditRequests {
     return null;
   };
 
-  public sendDocumentTransaction = () => Promise.resolve(null);
+  public sendDocumentTransaction: EditRequests["sendDocumentTransaction"] =
+    async (request) => {
+      await this.initialized.promise;
+      await this.rpc.proxy.request.documentTransaction(request);
+      return null;
+    };
 
   public addPackage: EditRequests["addPackage"] = async (request) => {
     return this.rpc.proxy.request.addPackage(request);
@@ -719,12 +754,6 @@ export class PyodideBridge implements RunRequests, EditRequests {
   }
 }
 
-export function createPyodideConnection(): IConnectionTransport {
-  return BasicTransport.withProducerCallback((callback) => {
-    PyodideBridge.INSTANCE.attachMessageConsumer(callback);
-  });
-}
-
 // Compose the worker name. The optional "::controller" suffix tells
 // getController.ts that the host page provides a custom /wasm/controller.js
 // and that the dynamic import should be attempted. Hosts opt in by setting

package/src/core/wasm/rpc.ts

@@ -9,19 +9,31 @@ import { Logger } from "@/utils/Logger";
 import type { WasmRuntimeConfig } from "./runtime-config";
 import { TRANSPORT_ID } from "./worker/constants";
 
+export interface PersistNotebookRequest {
+  filename: string;
+  contents: string;
+}
+
 export type ParentSchema = RPCSchema<{
+  requests: {
+    persistNotebook: { params: PersistNotebookRequest; response: null };
+  };
   messages: {
-    consumerReady: {};
     bootstrap: WasmRuntimeConfig;
   };
 }>;
 
-export function getWorkerRPC<WorkerSchema extends RPCSchema>(worker: Worker) {
+export function getWorkerRPC<WorkerSchema extends RPCSchema>(
+  worker: Worker,
+  requestHandler?: RPCOptions<ParentSchema, WorkerSchema>["requestHandler"],
+  maxRequestTime = 20_000,
+) {
   return createRPC<ParentSchema, WorkerSchema>({
     transport: createWorkerTransport(worker, {
       transportId: TRANSPORT_ID,
     }),
-    maxRequestTime: 20_000, // 20 seconds
+    requestHandler,
+    maxRequestTime,
     _debugHooks: {
       onSend: (message) => {
         Logger.debug("[rpc] Parent -> Worker", message);

package/src/core/wasm/store.ts

@@ -41,7 +41,7 @@ const urlFileStore: FileStore = {
 };
 
 export const domElementFileStore: FileStore = {
-  saveFile(contents: string) {
+  saveFile(_contents: string) {
     // Do nothing
   },
   readFile() {
@@ -54,7 +54,7 @@ export const domElementFileStore: FileStore = {
 };
 
 const remoteDefaultFileStore: FileStore = {
-  saveFile(contents: string) {
+  saveFile(_contents: string) {
     // Do nothing
   },
   readFile() {
@@ -79,7 +79,7 @@ export const mountConfigFileStore: FileStore = {
 };
 
 const emptyFileStore: FileStore = {
-  saveFile(contents: string) {
+  saveFile(_contents: string) {
     // Do nothing
   },
   readFile() {