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

package/src/core/host/wasm.ts

@@ -15,6 +15,7 @@ export interface WasmNotebookBridge {
   subscribe(
     cursor: string | null,
     emit: (event: unknown) => void,
+    signal: AbortSignal,
   ): Promise<() => void>;
   /** Deliver a call to the worker and resolve with its status and body. */
   call(call: unstable_NotebookCall): Promise<unknown>;
@@ -65,11 +66,15 @@ export function createWasmNotebookBackend(
         await abortable(
           () =>
             bridge
-              .subscribe(cursor, (event) => {
-                if (!signal.aborted) {
-                  message(event);
-                }
-              })
+              .subscribe(
+                cursor,
+                (event) => {
+                  if (!signal.aborted) {
+                    message(event);
+                  }
+                },
+                signal,
+              )
               .then((close) => {
                 unsubscribe = close;
                 if (signal.aborted) {

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

@@ -2,6 +2,7 @@
 
 import type { PyodideInterface } from "pyodide";
 import { describe, expect, it, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
 import { ReadonlyWasmController } from "../controller";
 
 class TestController extends ReadonlyWasmController {
@@ -50,14 +51,71 @@ function createPyodideStub(
   };
 }
 
-function startSession(controller: TestController, code: string) {
-  return controller.startSession({
+function prepareSession(controller: TestController, code: string) {
+  return controller.prepareSession({
     code,
     filename: `${code}.py`,
   });
 }
 
+async function startSession(controller: TestController, code: string) {
+  const session = await prepareSession(controller, code);
+  void session.start();
+  return session;
+}
+
 describe("WASM controller session lifecycle", () => {
+  it("releases a prepared session and ignores start after stop", async () => {
+    const { init, pyodide } = createPyodideStub();
+    const controller = new TestController();
+    controller.setPyodide(pyodide);
+
+    const session = await prepareSession(controller, "current");
+    await session.stop();
+    expect(init.destroy).toHaveBeenCalledOnce();
+    await session.start();
+    expect(init).not.toHaveBeenCalled();
+  });
+
+  it.each(["resolve", "reject"] as const)(
+    "releases a stopped session before its dependencies %s",
+    async (outcome) => {
+      const dependencies = new Deferred<void>();
+      const { init, loadPackagesFromImports, pyodide, bridge, stop } =
+        createPyodideStub();
+      loadPackagesFromImports.mockReturnValue(dependencies.promise);
+      const controller = new TestController();
+      controller.setPyodide(pyodide);
+      const session = await prepareSession(controller, "current");
+      const starting = session.start();
+      expect(init).not.toHaveBeenCalled();
+      await session.stop();
+      expect(init.destroy).toHaveBeenCalledOnce();
+      expect(bridge.destroy).toHaveBeenCalledOnce();
+      if (outcome === "resolve") {
+        dependencies.resolve();
+      } else {
+        dependencies.reject(new Error("Loading failed after stop"));
+      }
+      await starting;
+      await session.stop();
+      expect(stop).toHaveBeenCalledOnce();
+      expect(init).not.toHaveBeenCalled();
+    },
+  );
+
+  it("reports dependency failures when starting an active session", async () => {
+    const { init, loadPackagesFromImports, pyodide } = createPyodideStub();
+    loadPackagesFromImports.mockRejectedValue(new Error("Dependency failed"));
+    const controller = new TestController();
+    controller.setPyodide(pyodide);
+    const session = await prepareSession(controller, "current");
+    await expect(session.start()).rejects.toThrow("Dependency failed");
+    expect(init).not.toHaveBeenCalled();
+    expect(init.destroy).toHaveBeenCalledOnce();
+    await session.stop();
+  });
+
   it("stops and releases the active session", async () => {
     const { init, packages, pyodide, sessionResources, stop } =
       createPyodideStub();
@@ -115,7 +173,6 @@ describe("WASM controller session lifecycle", () => {
 
     await startSession(controller, "first");
     await startSession(controller, "second");
-    second.bridge.destroy();
 
     expect(first.stop.destroy).not.toHaveBeenCalled();
 
@@ -125,6 +182,8 @@ describe("WASM controller session lifecycle", () => {
     expect(second.stop).toHaveBeenCalledOnce();
     expect(first.stop.destroy).toHaveBeenCalledOnce();
     expect(second.stop.destroy).toHaveBeenCalledOnce();
+    expect(first.bridge.destroy).toHaveBeenCalledOnce();
+    expect(second.bridge.destroy).toHaveBeenCalledOnce();
   });
 
   it("continues stopping sessions after one stop fails", async () => {

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

@@ -3,7 +3,7 @@ import type { PyodideInterface } from "pyodide";
 import { defaultUserConfig } from "@/core/config/config-schema";
 import { DefaultWasmController } from "@/core/wasm/worker/bootstrap";
 import { WasmFileSystem } from "@/core/wasm/worker/fs";
-import type { SerializedBridge } from "@/core/wasm/worker/types";
+import type { WasmSession } from "@/core/wasm/worker/types";
 import { Logger } from "../../../utils/Logger";
 
 export class ReadonlyWasmController extends DefaultWasmController {
@@ -35,17 +35,15 @@ export class ReadonlyWasmController extends DefaultWasmController {
     return { code, filename };
   }
 
-  public override async startSession(opts: {
+  public override async prepareSession(opts: {
     code: string;
     filename: string | null;
-  }): Promise<SerializedBridge> {
-    const bridge = super.startSession({
+  }): Promise<WasmSession> {
+    return super.prepareSession({
       queryParameters: {},
       code: opts.code,
       filename: opts.filename,
       userConfig: defaultUserConfig(),
     });
-
-    return bridge;
   }
 }

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

@@ -1,7 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import type { PyodideInterface } from "pyodide";
-import type { PyProxy } from "pyodide/ffi";
 import {
   createRPC,
   createRPCRequestHandler,
@@ -56,8 +55,6 @@ 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(
@@ -75,15 +72,22 @@ async function startSession(
       code: opts.code,
       filename: `app-${opts.appId}.py`,
     });
-    const nextBridge = await self.controller.startSession(notebook);
-    sessions.add(nextBridge);
+    const session = await self.controller.prepareSession(notebook);
     // Follow before the session instantiates. Each subscription begins with
     // its own kernel-ready and the session's retained notifications.
-    await follow(nextBridge, opts.sessionGeneration);
+    try {
+      await follow(session.bridge, opts.sessionGeneration);
+    } catch (error) {
+      await session.stop();
+      throw error;
+    }
+    void session.start().catch((error: unknown) => {
+      rpc.send.initializedError({ error: prettyError(error) });
+    });
     if (!activeSession) {
       activeSession = {
         appId: opts.appId,
-        bridge: nextBridge,
+        bridge: session.bridge,
         sessionGeneration: opts.sessionGeneration,
       };
     }
@@ -100,10 +104,6 @@ async function stopActiveSession(): Promise<void> {
   // Stopping closes every session and its notebook subscriptions.
   await self.controller.stopSession();
   activeSession = undefined;
-  for (const bridge of sessions) {
-    (bridge as unknown as PyProxy).destroy();
-  }
-  sessions.clear();
 }
 
 async function follow(

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

@@ -11,6 +11,9 @@ const {
   mockSaveNotebook,
   mockBootstrap,
   rpcListeners,
+  mockNotebookInfo,
+  mockNotebookSubscribe,
+  mockNotebookUnsubscribe,
 } = vi.hoisted(() => ({
   mockBridge: vi.fn(),
   mockNotebookReadFile: vi.fn(),
@@ -18,6 +21,9 @@ const {
   mockSaveNotebook: vi.fn(),
   mockBootstrap: vi.fn(),
   rpcListeners: {} as Record<string, () => void>,
+  mockNotebookInfo: vi.fn(),
+  mockNotebookSubscribe: vi.fn(),
+  mockNotebookUnsubscribe: vi.fn(),
 }));
 
 // Mock browser globals before any imports
@@ -50,6 +56,9 @@ vi.mock("@/core/wasm/rpc", () => ({
         readFile: vi.fn(),
         readNotebook: mockReadNotebook,
         saveNotebook: mockSaveNotebook,
+        notebookInfo: mockNotebookInfo,
+        notebookSubscribe: mockNotebookSubscribe,
+        notebookUnsubscribe: mockNotebookUnsubscribe,
       },
     },
     addMessageListener: (event: string, cb: () => void) => {
@@ -217,6 +226,39 @@ describe("PyodideBridge.sendSave", () => {
   });
 });
 
+it("cancels a pending notebook subscription before its RPC response", async () => {
+  vi.clearAllMocks();
+  store.set(initialModeAtom, "read");
+  rpcListeners.initialized();
+  const pending = new Deferred<void>();
+  mockNotebookInfo.mockResolvedValue({ id: "n1" });
+  mockNotebookSubscribe.mockReturnValue(pending.promise);
+  mockNotebookUnsubscribe.mockResolvedValue(undefined);
+  const controller = new AbortController();
+  try {
+    const backend = await PyodideBridge.INSTANCE.createNotebookBackend();
+    const following = backend.transport(
+      { cursor: null, signal: controller.signal },
+      { open: vi.fn(), message: vi.fn() },
+    );
+    await vi.waitFor(() =>
+      expect(mockNotebookSubscribe).toHaveBeenCalledOnce(),
+    );
+    const [{ id }] = mockNotebookSubscribe.mock.calls[0];
+    controller.abort();
+    await following;
+    expect(mockNotebookUnsubscribe.mock.calls).toEqual([[id]]);
+
+    pending.resolve();
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(mockNotebookUnsubscribe.mock.calls).toEqual([[id]]);
+  } finally {
+    controller.abort();
+    pending.resolve();
+    store.set(initialModeAtom, undefined);
+  }
+});
+
 describe("getWasmWorkerName", () => {
   afterEach(() => {
     delete (window as unknown as { __MARIMO_HAS_WASM_CONTROLLER__?: boolean })

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

@@ -0,0 +1,182 @@
+/* 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),
+    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, event, pyodide } = await startWorker();
+  try {
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(rpc.send.initialized).toHaveBeenCalledOnce();
+    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();
+  } 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

@@ -8,6 +8,7 @@ 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";
@@ -211,16 +212,19 @@ export class PyodideBridge implements RunRequests, EditRequests {
     const { id } = await this.rpc.proxy.request.notebookInfo();
     return createWasmNotebookBackend({
       notebookId: id,
-      subscribe: async (cursor, emit) => {
+      subscribe: async (cursor, emit, signal) => {
         await this.initialized.promise;
+        signal.throwIfAborted();
         const id = generateUUID();
         this.notebookConsumers.set(id, emit);
-        const close = () => {
+        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;

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

@@ -1,13 +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 { 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;
@@ -43,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");
@@ -136,18 +136,15 @@ export class DefaultWasmController implements WasmController {
     });
   }
 
-  public async startSession(opts: {
-    queryParameters: Record<string, string | string[]>;
-    code: string;
-    filename: string | null;
-    userConfig: UserConfig;
-  }): Promise<SerializedBridge> {
+  public async prepareSession(
+    opts: Parameters<WasmController["prepareSession"]>[0],
+  ): Promise<WasmSession> {
     const sessionGeneration = this.sessionGeneration;
     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(
       `
@@ -222,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/types.ts

@@ -54,18 +54,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;
-  }): 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 {