@marimo-team/islands 0.25.3-dev16 → 0.25.3-dev167

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 "../transports/wasm";
 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 { activeNotebookReplicaAtom } from "../session/active-notebook";
 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(activeNotebookReplicaAtom);
+    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() {

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

@@ -1,15 +1,13 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 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 { once } from "../../../utils/once";
 import { Logger } from "../../../utils/Logger";
 import { WasmFileSystem } from "./fs";
 import { getMarimoWheel } from "./getMarimoWheel";
 import { t } from "./tracer";
-import type { SerializedBridge, WasmController } from "./types";
+import type { SerializedBridge, WasmController, WasmSession } from "./types";
 import { shouldLoadDuckDBPackages } from "../utils";
 
 const MAKE_SNAPSHOT = false;
@@ -45,7 +43,7 @@ export class DefaultWasmController implements WasmController {
   protected pyodide: PyodideInterface | null = null;
   private packageLoadQueue = Promise.resolve();
   private sessionGeneration = 0;
-  private activeSessionStops = new Set<PyCallable>();
+  private sessions = new Set<WasmSession>();
 
   public get requirePyodide() {
     invariant(this.pyodide, "Pyodide not loaded");
@@ -85,7 +83,7 @@ export class DefaultWasmController implements WasmController {
         packages: [
           "micropip",
           "msgspec",
-          getMarimoWheel(opts.version),
+          getMarimoWheel(),
           "Markdown",
           "pymdown-extensions",
           "narwhals",
@@ -138,28 +136,15 @@ export class DefaultWasmController implements WasmController {
     });
   }
 
-  public async startSession(opts: {
-    queryParameters: Record<string, string | string[]>;
-    code: string;
-    filename: string | null;
-    onMessage: (message: JsonString<NotificationPayload>) => void;
-    userConfig: UserConfig;
-  }): Promise<SerializedBridge> {
+  public async prepareSession(
+    opts: Parameters<WasmController["prepareSession"]>[0],
+  ): Promise<WasmSession> {
     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;
 
-    const span = t.startSpan("startSession.runPython");
+    const span = t.startSpan("prepareSession.runPython");
     const nbFilename = filename || WasmFileSystem.NOTEBOOK_FILENAME;
     const sessionResources = this.requirePyodide.runPython(
       `
@@ -169,14 +154,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
@@ -198,6 +181,7 @@ export class DefaultWasmController implements WasmController {
           if task is not None:
             await asyncio.gather(task, return_exceptions=True)
           if bridge is not None:
+            bridge.close()
             bridge.session = None
           task = None
           kernel_task = None
@@ -235,45 +219,65 @@ export class DefaultWasmController implements WasmController {
     } finally {
       packagesProxy.destroy();
     }
-    this.activeSessionStops.add(stopSession);
-
-    // Fire and forget:
-    // Load notebook dependencies and instantiate the session
-    // We don't want to wait for this to finish,
-    // so we can show the initial code immediately giving
-    // a sense of responsiveness.
+    let stopped = false;
+    let stopping: Promise<void> | undefined;
+    const releaseInit = once(() => initSession.destroy());
+    const isStopped = () =>
+      stopped || sessionGeneration !== this.sessionGeneration;
+    // Load in parallel with subscription setup; only start() executes cells.
     const dependenciesReady = this.packageLoadQueue.then(() => {
-      if (sessionGeneration !== this.sessionGeneration) {
+      if (isStopped()) {
         return;
       }
       return this.loadNotebookDeps(code, foundPackages);
     });
     this.packageLoadQueue = dependenciesReady.catch(() => undefined);
-    void dependenciesReady
-      .then(() => {
-        if (sessionGeneration !== this.sessionGeneration) {
-          return;
+    const session: WasmSession = {
+      bridge: bridgeProxy as unknown as SerializedBridge,
+      start: once(async () => {
+        try {
+          if (isStopped()) {
+            return;
+          }
+          await dependenciesReady;
+          if (!isStopped()) {
+            initSession(userConfig.runtime.auto_instantiate);
+          }
+        } catch (error) {
+          if (!isStopped()) {
+            throw error;
+          }
+        } finally {
+          releaseInit();
         }
-        return initSession(userConfig.runtime.auto_instantiate);
-      })
-      .catch((error: unknown) => {
-        Logger.error("Failed to load notebook dependencies", error);
-      })
-      .finally(() => initSession.destroy());
-
-    return bridgeProxy as unknown as SerializedBridge;
+      }),
+      stop: () => {
+        stopped = true;
+        releaseInit();
+        stopping ??= (async () => {
+          await stopSession();
+          stopSession.destroy();
+          bridgeProxy.destroy();
+          this.sessions.delete(session);
+        })().catch((error: unknown) => {
+          stopping = undefined;
+          throw error;
+        });
+        return stopping;
+      },
+    };
+    this.sessions.add(session);
+    return session;
   }
 
   public async stopSession(): Promise<void> {
     this.sessionGeneration += 1;
-    const stops = [...this.activeSessionStops];
+    const sessions = [...this.sessions];
     let hasFailure = false;
     let firstFailure: unknown;
-    for (const stop of stops) {
+    for (const session of sessions) {
       try {
-        await stop();
-        this.activeSessionStops.delete(stop);
-        stop.destroy();
+        await session.stop();
       } catch (error) {
         if (!hasFailure) {
           hasFailure = true;

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

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-export function getMarimoWheel(version: string) {
+export function getMarimoWheel() {
   if (import.meta.env.DEV && import.meta.env.VITE_WASM_MARIMO_PREBUILT_WHEEL) {
     return "marimo-base";
   }

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

@@ -7,12 +7,11 @@ import {
   createWorkerParentTransport,
   type RPCSchema,
 } from "rpc-anywhere";
-import type { SaveNotebookRequest } from "@/core/network/types";
 import { decodeUtf8 } from "@/utils/strings";
 import { Deferred } from "../../../utils/Deferred";
 import { prettyError } from "../../../utils/errors";
 import { Logger } from "../../../utils/Logger";
-import type { ParentSchema } from "../rpc";
+import type { ParentSchema, PersistNotebookRequest } from "../rpc";
 import type { WasmRuntimeConfig } from "../runtime-config";
 import { TRANSPORT_ID } from "./constants";
 import { WasmFileSystem } from "./fs";
@@ -70,14 +69,14 @@ const requestHandler = createRPCRequestHandler({
     await pyodideReadyPromise; // Make sure loading is done
     return WasmFileSystem.readNotebook(self.pyodide);
   },
-  saveNotebook: async (opts: SaveNotebookRequest) => {
-    await pyodideReadyPromise; // Make sure loading is done
-    const saveFile = self.pyodide.runPython(`
-      from marimo._pyodide.bootstrap import save_file
-
-      save_file
-    `);
-    await saveFile(JSON.stringify(opts), WasmFileSystem.NOTEBOOK_FILENAME);
+  writeNotebook: async ({ filename, contents }: PersistNotebookRequest) => {
+    await pyodideReadyPromise;
+    const directory = filename.substring(0, filename.lastIndexOf("/"));
+    if (directory) {
+      self.pyodide.FS.mkdirTree(directory);
+    }
+    self.pyodide.FS.writeFile(filename, contents);
+    self.pyodide.FS.writeFile(WasmFileSystem.NOTEBOOK_FILENAME, contents);
     await WasmFileSystem.persistFilesToRemote(self.pyodide);
   },
 });

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,
@@ -34,7 +30,6 @@ import type {
   FormatCellsRequest,
   FormatResponse,
   SaveAppConfigurationRequest,
-  SaveNotebookRequest,
   SaveUserConfigurationRequest,
   Snippets,
 } from "../../network/types";
@@ -58,19 +53,21 @@ export interface WasmController {
     code: string;
     filename: string | null;
   }): Promise<{ code: string; filename: string }>;
-  /**
-   * Start the session
-   * @param opts.code - The code to start with
-   * @param opts.fallbackCode - The code to fallback to
-   * @param opts.filename - The filename to start with
-   */
-  startSession(opts: {
+  /** Prepare a session and load its dependencies without executing cells. */
+  prepareSession(opts: {
     queryParameters: Record<string, string | string[]>;
     code: string;
     filename: string | null;
     userConfig: UserConfig;
-    onMessage: (message: JsonString<NotificationPayload>) => void;
-  }): Promise<SerializedBridge>;
+  }): Promise<WasmSession>;
+}
+
+export interface WasmSession {
+  readonly bridge: SerializedBridge;
+  /** Start once dependencies are ready. Repeated calls share the same start. */
+  start(): Promise<void>;
+  /** Prevent a pending start and release the session and its bridge. */
+  stop(): Promise<void>;
 }
 
 export interface RawBridge {
@@ -81,7 +78,6 @@ export interface RawBridge {
   read_snippets(): Promise<Snippets>;
   get_environment_info(): Promise<EnvironmentInfo>;
   format(request: FormatCellsRequest): Promise<FormatResponse>;
-  save(request: SaveNotebookRequest): Promise<string>;
   copy(request: CopyNotebookRequest): Promise<string>;
   save_app_config(request: SaveAppConfigurationRequest): Promise<string>;
   save_user_config(request: SaveUserConfigurationRequest): Promise<null>;
@@ -114,11 +110,26 @@ export type BridgePayload<T extends keyof RawBridge> = T extends keyof RawBridge
   : undefined;
 
 export type SerializedBridge = {
-  [P in Exclude<keyof RawBridge, "save">]: RawBridge[P] extends (
+  [P in keyof RawBridge]: RawBridge[P] extends (
     payload: string,
   ) => Promise<unknown>
     ? (payload: string) => Promise<string>
     : RawBridge[P];
 } & {
-  save(payload: string): Promise<void>;
+  notebook: {
+    readonly id: string;
+    subscribe(
+      id: string,
+      cursor: string | null,
+      emit: (event: string) => void,
+    ): Promise<void>;
+    unsubscribe(id: string): void;
+    call(call: string): Promise<string>;
+    transaction(request: string): Promise<void>;
+    set_layout(layout: string): Promise<void>;
+    set_persistence(
+      persist: (contents: string, filename: string) => Promise<void>,
+    ): void;
+    destroy(): void;
+  };
 };