@marimo-team/islands 0.23.15-dev30 → 0.23.15-dev32

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

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import type { PyodideInterface } from "pyodide";
+import type { PyProxy } from "pyodide/ffi";
 import {
   createRPC,
   createRPCRequestHandler,
@@ -16,7 +17,6 @@ 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 { Deferred } from "../../../utils/Deferred";
 import { prettyError } from "../../../utils/errors";
 import { invariant } from "../../../utils/invariant";
 import { ReadonlyWasmController } from "./controller";
@@ -50,59 +50,128 @@ const messageBuffer = new MessageBuffer(
     rpc.send.kernelMessage({ message });
   },
 );
-const bridgeReady = new Deferred<SerializedBridge>();
+interface SessionRequest {
+  code: string;
+  appId: string;
+  sessionGeneration: number;
+}
+
+let activeSession:
+  | (Omit<SessionRequest, "code"> & { bridge: SerializedBridge })
+  | undefined;
+let sessionQueue = Promise.resolve();
+
+async function startSession(
+  opts: SessionRequest,
+  replace: boolean,
+): Promise<void> {
+  await pyodideReadyPromise;
+
+  try {
+    invariant(self.controller, "Controller not loaded");
+    if (replace) {
+      await stopActiveSession();
+    }
+    const notebook = await self.controller.mountFilesystem({
+      code: opts.code,
+      filename: `app-${opts.appId}.py`,
+    });
+    const nextBridge = await self.controller.startSession({
+      ...notebook,
+      onMessage: messageBuffer.push,
+    });
+    if (activeSession) {
+      (nextBridge as unknown as PyProxy).destroy();
+    } else {
+      activeSession = {
+        appId: opts.appId,
+        bridge: nextBridge,
+        sessionGeneration: opts.sessionGeneration,
+      };
+    }
+    rpc.send.initialized({});
+  } catch (error) {
+    rpc.send.initializedError({
+      error: prettyError(error),
+    });
+    throw error;
+  }
+}
+
+async function stopActiveSession(): Promise<void> {
+  const session = activeSession;
+  await self.controller.stopSession();
+  activeSession = undefined;
+  (session?.bridge as unknown as PyProxy | undefined)?.destroy();
+}
+
+function enqueueSession<T>(operation: () => Promise<T>): Promise<T> {
+  const result = sessionQueue.then(operation);
+  sessionQueue = result.then(
+    () => undefined,
+    () => undefined,
+  );
+  return result;
+}
 
 // Handle RPC requests
 const requestHandler = createRPCRequestHandler({
   /**
    * Start the session
    */
-  startSession: async (opts: { code: string; appId: string }) => {
-    await pyodideReadyPromise; // Make sure loading is done
+  startSession: async (opts: SessionRequest) => {
+    await enqueueSession(() => startSession(opts, false));
+  },
 
-    try {
+  replaceSession: async (opts: SessionRequest) => {
+    await enqueueSession(() => startSession(opts, true));
+  },
+
+  stopSession: async (opts: { appId: string; sessionGeneration: number }) => {
+    await enqueueSession(async () => {
+      if (
+        activeSession?.appId !== opts.appId ||
+        activeSession?.sessionGeneration !== opts.sessionGeneration
+      ) {
+        return;
+      }
       invariant(self.controller, "Controller not loaded");
-      const notebook = await self.controller.mountFilesystem({
-        code: opts.code,
-        filename: `app-${opts.appId}.py`,
-      });
-      const bridge = await self.controller.startSession({
-        ...notebook,
-        onMessage: messageBuffer.push,
-      });
-      bridgeReady.resolve(bridge);
-      rpc.send.initialized({});
-    } catch (error) {
-      rpc.send.initializedError({
-        error: prettyError(error),
-      });
-    }
-    return;
+      await stopActiveSession();
+    });
   },
 
   /**
    * Load packages
    */
-  loadPackages: async (code: string) => {
+  loadPackages: async (opts: {
+    appId: string;
+    code: string;
+    sessionGeneration: number;
+  }) => {
     await pyodideReadyPromise; // Make sure loading is done
+    await enqueueSession(async () => {
+      requireActiveBridge(opts);
+
+      let { code } = opts;
 
-    if (shouldLoadDuckDBPackages(code)) {
-      // Add pandas and duckdb to the code for mo.sql and for remote duckdb sources
-      code = `import pandas\n${code}`;
-      code = `import duckdb\n${code}`;
-      code = `import sqlglot\n${code}`;
-
-      // Polars + SQL requires pyarrow, and installing
-      // after notebook load does not work. As a heuristic,
-      // if it appears that the notebook uses polars, add pyarrow.
-      if (code.includes("polars")) {
-        code = `import pyarrow\n${code}`;
+      if (shouldLoadDuckDBPackages(code)) {
+        // Add pandas and duckdb to the code for mo.sql and for remote duckdb sources
+        code = `import pandas\n${code}`;
+        code = `import duckdb\n${code}`;
+        code = `import sqlglot\n${code}`;
+
+        // Polars + SQL requires pyarrow, and installing
+        // after notebook load does not work. As a heuristic,
+        // if it appears that the notebook uses polars, add pyarrow.
+        if (code.includes("polars")) {
+          code = `import pyarrow\n${code}`;
+        }
       }
-    }
 
-    await self.pyodide.loadPackagesFromImports(code, {
-      messageCallback: Logger.log,
-      errorCallback: Logger.error,
+      await self.pyodide.loadPackagesFromImports(code, {
+        messageCallback: Logger.log,
+        errorCallback: Logger.error,
+      });
     });
   },
 
@@ -110,37 +179,56 @@ const requestHandler = createRPCRequestHandler({
    * Call a function on the bridge
    */
   bridge: async (opts: {
+    appId: string;
     functionName: keyof RawBridge;
     payload: {} | undefined | null;
+    sessionGeneration: number;
   }) => {
     await pyodideReadyPromise; // Make sure loading is done
 
     const { functionName, payload } = opts;
 
     // Perform the function call to the Python bridge
-    const bridge = await bridgeReady.promise;
-
-    // Serialize the payload
-    const payloadString =
-      payload == null
-        ? null
-        : typeof payload === "string"
-          ? payload
-          : JSON.stringify(payload);
-
-    // Make the request
-    const response =
-      payloadString == null
-        ? // @ts-expect-error ehh TypeScript
-          await bridge[functionName]()
-        : // @ts-expect-error ehh TypeScript
-          await bridge[functionName](payloadString);
-
-    // Post the response back to the main thread
-    return typeof response === "string" ? JSON.parse(response) : response;
+    return enqueueSession(async () => {
+      const bridge = requireActiveBridge(opts);
+
+      // Serialize the payload
+      const payloadString =
+        payload == null
+          ? null
+          : typeof payload === "string"
+            ? payload
+            : JSON.stringify(payload);
+
+      // Make the request
+      const response =
+        payloadString == null
+          ? // @ts-expect-error ehh TypeScript
+            await bridge[functionName]()
+          : // @ts-expect-error ehh TypeScript
+            await bridge[functionName](payloadString);
+
+      // Post the response back to the main thread
+      return typeof response === "string" ? JSON.parse(response) : response;
+    });
   },
 });
 
+function requireActiveBridge(opts: {
+  appId: string;
+  sessionGeneration: number;
+}): SerializedBridge {
+  const session = activeSession;
+  if (
+    !session ||
+    opts.appId !== session.appId ||
+    opts.sessionGeneration !== session.sessionGeneration
+  ) {
+    throw new Error("The marimo app session is no longer active.");
+  }
+  return session.bridge;
+}
+
 // create the iframe's schema
 export type WorkerSchema = RPCSchema<
   {

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

@@ -1,5 +1,6 @@
 /* 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";
@@ -12,6 +13,12 @@ import type { SerializedBridge, WasmController } from "./types";
 import { shouldLoadDuckDBPackages } from "../utils";
 
 const MAKE_SNAPSHOT = false;
+type SessionResources = [
+  bridge: PyProxy,
+  init: PyCallable,
+  packages: PyProxy,
+  stop: PyCallable,
+];
 
 // This class initializes the wasm environment
 // We would like this initialization to be parallelizable
@@ -25,6 +32,9 @@ const MAKE_SNAPSHOT = false;
 
 export class DefaultWasmController implements WasmController {
   protected pyodide: PyodideInterface | null = null;
+  private packageLoadQueue = Promise.resolve();
+  private sessionGeneration = 0;
+  private activeSessionStops = new Set<PyCallable>();
 
   get requirePyodide() {
     invariant(this.pyodide, "Pyodide not loaded");
@@ -111,6 +121,7 @@ export class DefaultWasmController implements WasmController {
     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
@@ -126,47 +137,129 @@ export class DefaultWasmController implements WasmController {
 
     const span = t.startSpan("startSession.runPython");
     const nbFilename = filename || WasmFileSystem.NOTEBOOK_FILENAME;
-    const [bridge, init, packages] = this.requirePyodide.runPython(
+    const sessionResources = this.requirePyodide.runPython(
       `
       print("[py] Starting marimo...")
       import asyncio
+      import gc
       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"
 
-      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(),
-      )
+      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
 
-      def init(auto_instantiate=True):
-        instantiate(session, auto_instantiate)
-        asyncio.create_task(session.start())
+        def init(auto_instantiate=True):
+          nonlocal session_task
+          instantiate(session, auto_instantiate)
+          session_task = asyncio.create_task(session.start())
 
-      # Find the packages to install
-      with open("${nbFilename}", "r") as f:
-        packages = session.find_packages(f.read())
+        async def stop():
+          nonlocal bridge, session, session_task
+          task = session_task
+          kernel_task = getattr(session, "kernel_task", None)
+          if kernel_task is None:
+            if task is not None:
+              task.cancel()
+          else:
+            kernel_task.stop()
+          if task is not None:
+            await asyncio.gather(task, return_exceptions=True)
+          if bridge is not None:
+            bridge.session = None
+          task = None
+          kernel_task = None
+          session_task = None
+          session = None
+          bridge = None
+          gc.collect()
 
-      bridge, init, packages`,
-    );
-    span.end();
+        with open("${nbFilename}", "r") as f:
+          packages = session.find_packages(f.read())
 
-    const foundPackages = new Set<string>(packages.toJs());
+        return bridge, init, packages, stop
+
+      create_session_resources()`,
+    ) as PyProxy;
+    span.end();
+    let bridgeProxy!: PyProxy;
+    let initSession!: PyCallable;
+    let packagesProxy!: PyProxy;
+    let stopSession!: PyCallable;
+    try {
+      [bridgeProxy, initSession, packagesProxy, stopSession] =
+        sessionResources as unknown as SessionResources;
+    } finally {
+      sessionResources.destroy();
+    }
+    let foundPackages: Set<string>;
+    try {
+      foundPackages = new Set<string>(packagesProxy.toJs());
+    } catch (error) {
+      bridgeProxy.destroy();
+      initSession.destroy();
+      stopSession.destroy();
+      throw error;
+    } 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.
-    void this.loadNotebookDeps(code, foundPackages).then(() => {
-      return init(userConfig.runtime.auto_instantiate);
+    const dependenciesReady = this.packageLoadQueue.then(() => {
+      if (sessionGeneration !== this.sessionGeneration) {
+        return;
+      }
+      return this.loadNotebookDeps(code, foundPackages);
     });
+    this.packageLoadQueue = dependenciesReady.catch(() => undefined);
+    void dependenciesReady
+      .then(() => {
+        if (sessionGeneration !== this.sessionGeneration) {
+          return;
+        }
+        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;
+  }
 
-    return bridge;
+  async stopSession(): Promise<void> {
+    this.sessionGeneration += 1;
+    const stops = [...this.activeSessionStops];
+    let hasFailure = false;
+    let firstFailure: unknown;
+    for (const stop of stops) {
+      try {
+        await stop();
+        this.activeSessionStops.delete(stop);
+        stop.destroy();
+      } catch (error) {
+        if (!hasFailure) {
+          hasFailure = true;
+          firstFailure = error;
+        }
+      }
+    }
+    if (hasFailure) {
+      throw firstFailure;
+    }
   }
 
   private async loadNotebookDeps(code: string, foundPackages: Set<string>) {

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

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