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

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;
+  };
 };

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,12 @@ 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,
+  NotebookDocumentTransactionRequest,
   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 +26,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 +70,66 @@ 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({
+  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 +153,34 @@ 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);
+      const notebook = session.bridge.notebook;
+      try {
+        notebook.set_persistence(async (contents, filename) => {
+          // Keep file operations in the execution worker from syncing stale source.
+          self.pyodide.FS.writeFile(filename, contents);
+          self.pyodide.FS.writeFile(WasmFileSystem.NOTEBOOK_FILENAME, contents);
+          await rpc.proxy.request.persistNotebook({ filename, contents });
+        });
+      } finally {
+        notebook.destroy();
+      }
+      sessionReady.resolve(session);
+      rpc.send.initializingMessage({
+        message: "Initializing notebook...",
+      });
+      rpc.send.initialized({});
       workerInitSpan.end("ok");
     } catch (error) {
       rpc.send.initializedError({
@@ -249,13 +302,21 @@ const requestHandler = createRPCRequestHandler({
     };
   },
 
-  /**
-   * Save the notebook
-   */
-  saveNotebook: async (opts: SaveNotebookRequest) => {
-    await pyodideReadyPromise; // Make sure loading is done
-    const bridge = await bridgeReady.promise;
-    await bridge.save(JSON.stringify(opts));
+  documentTransaction: async (request: NotebookDocumentTransactionRequest) => {
+    const notebook = await getNotebook();
+    try {
+      await notebook.transaction(JSON.stringify(request));
+    } finally {
+      notebook.destroy();
+    }
+  },
+  setNotebookLayout: async (layout: SaveNotebookRequest["layout"]) => {
+    const notebook = await getNotebook();
+    try {
+      await notebook.set_layout(JSON.stringify(layout));
+    } finally {
+      notebook.destroy();
+    }
   },
 
   /**
@@ -302,7 +363,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 +396,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
@@ -355,22 +415,15 @@ const rpc = createRPC<WorkerSchema, ParentSchema>({
     transportId: TRANSPORT_ID,
   }),
   requestHandler,
+  // Persistence can include bootstrapping the storage worker.
+  maxRequestTime: Infinity,
 });
 
 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",
   "rename_file",
   "create_file_or_directory",
   "delete_file_or_directory",