@marimo-team/islands 0.23.15-dev4 → 0.23.15-dev40

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

@@ -0,0 +1,173 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { PyodideInterface } from "pyodide";
+import { describe, expect, it, vi } from "vitest";
+import { ReadonlyWasmController } from "../controller";
+
+class TestController extends ReadonlyWasmController {
+  setPyodide(pyodide: PyodideInterface) {
+    this.pyodide = pyodide;
+  }
+}
+
+function createSessionResources(
+  stopImplementation: () => Promise<void> = () => Promise.resolve(),
+) {
+  const bridge = { destroy: vi.fn() };
+  const init = Object.assign(vi.fn(), { destroy: vi.fn() });
+  const stop = Object.assign(vi.fn(stopImplementation), {
+    destroy: vi.fn(),
+  });
+  const packages = { destroy: vi.fn(), toJs: () => [] };
+  const sessionResources = Object.assign([bridge, init, packages, stop], {
+    destroy: vi.fn(),
+  });
+
+  return { bridge, init, packages, sessionResources, stop };
+}
+
+function createPyodideStub(
+  sessions = Array.from({ length: 4 }, () => createSessionResources()),
+) {
+  const [first] = sessions;
+  if (!first) {
+    throw new Error("At least one session is required");
+  }
+  const loadPackagesFromImports = vi.fn().mockResolvedValue(undefined);
+  const pyodide = {
+    runPython: vi
+      .fn()
+      .mockImplementation(() => sessions.shift()?.sessionResources),
+    loadPackagesFromImports,
+    loadedPackages: {},
+    runPythonAsync: vi.fn(),
+  } as unknown as PyodideInterface;
+
+  return {
+    ...first,
+    loadPackagesFromImports,
+    pyodide,
+  };
+}
+
+function startSession(controller: TestController, code: string) {
+  return controller.startSession({
+    code,
+    filename: `${code}.py`,
+    onMessage: vi.fn(),
+  });
+}
+
+describe("WASM controller session lifecycle", () => {
+  it("stops and releases the active session", async () => {
+    const { init, packages, pyodide, sessionResources, stop } =
+      createPyodideStub();
+    const controller = new TestController();
+    controller.setPyodide(pyodide);
+
+    await startSession(controller, "current");
+    await vi.waitFor(() => expect(init).toHaveBeenCalledOnce());
+    await controller.stopSession();
+
+    expect(stop).toHaveBeenCalledOnce();
+    expect(stop.destroy).toHaveBeenCalledOnce();
+    expect(sessionResources.destroy).toHaveBeenCalledOnce();
+    expect(packages.destroy).toHaveBeenCalledOnce();
+    expect(init.destroy).toHaveBeenCalledOnce();
+  });
+
+  it("loads dependencies serially and skips superseded sessions", async () => {
+    let finishFirstLoad!: () => void;
+    const { loadPackagesFromImports, pyodide } = createPyodideStub();
+    loadPackagesFromImports
+      .mockReturnValueOnce(
+        new Promise<void>((resolve) => {
+          finishFirstLoad = resolve;
+        }),
+      )
+      .mockResolvedValueOnce(undefined);
+    const controller = new TestController();
+    controller.setPyodide(pyodide);
+
+    await startSession(controller, "first_dependency");
+    await startSession(controller, "superseded_dependency");
+    await controller.stopSession();
+    await startSession(controller, "current_dependency");
+
+    expect(loadPackagesFromImports).toHaveBeenCalledOnce();
+    finishFirstLoad();
+    await vi.waitFor(() =>
+      expect(loadPackagesFromImports).toHaveBeenCalledTimes(2),
+    );
+
+    const loadedSources = loadPackagesFromImports.mock.calls.map(
+      ([source]) => source as string,
+    );
+    expect(loadedSources[1]).toContain("current_dependency");
+    expect(loadedSources.join("\n")).not.toContain("superseded_dependency");
+  });
+
+  it("stops every session created for a multi-app page", async () => {
+    const first = createSessionResources();
+    const second = createSessionResources();
+    const { pyodide } = createPyodideStub([first, second]);
+    const controller = new TestController();
+    controller.setPyodide(pyodide);
+
+    await startSession(controller, "first");
+    await startSession(controller, "second");
+    second.bridge.destroy();
+
+    expect(first.stop.destroy).not.toHaveBeenCalled();
+
+    await controller.stopSession();
+
+    expect(first.stop).toHaveBeenCalledOnce();
+    expect(second.stop).toHaveBeenCalledOnce();
+    expect(first.stop.destroy).toHaveBeenCalledOnce();
+    expect(second.stop.destroy).toHaveBeenCalledOnce();
+  });
+
+  it("continues stopping sessions after one stop fails", async () => {
+    const failure = new Error("stop failed");
+    const first = createSessionResources(() => Promise.reject(failure));
+    const second = createSessionResources();
+    const { pyodide } = createPyodideStub([first, second]);
+    const controller = new TestController();
+    controller.setPyodide(pyodide);
+
+    await startSession(controller, "first");
+    await startSession(controller, "second");
+
+    await expect(controller.stopSession()).rejects.toThrow("stop failed");
+    expect(first.stop).toHaveBeenCalledOnce();
+    expect(second.stop).toHaveBeenCalledOnce();
+    expect(first.stop.destroy).not.toHaveBeenCalled();
+    expect(second.stop.destroy).toHaveBeenCalledOnce();
+
+    first.stop.mockResolvedValueOnce(undefined);
+    await controller.stopSession();
+
+    expect(first.stop).toHaveBeenCalledTimes(2);
+    expect(first.stop.destroy).toHaveBeenCalledOnce();
+    expect(second.stop).toHaveBeenCalledOnce();
+  });
+
+  it("reports and retries a falsy stop failure", async () => {
+    const session = createSessionResources(() => Promise.reject(false));
+    const { pyodide } = createPyodideStub([session]);
+    const controller = new TestController();
+    controller.setPyodide(pyodide);
+
+    await startSession(controller, "current");
+
+    await expect(controller.stopSession()).rejects.toBe(false);
+    expect(session.stop.destroy).not.toHaveBeenCalled();
+
+    session.stop.mockResolvedValueOnce(undefined);
+    await controller.stopSession();
+
+    expect(session.stop).toHaveBeenCalledTimes(2);
+    expect(session.stop.destroy).toHaveBeenCalledOnce();
+  });
+});

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/runtime/__tests__/runtime.test.ts

@@ -154,6 +154,70 @@ describe("RuntimeManager", () => {
     });
   });
 
+  describe("formatNavigableHttpURL", () => {
+    it("adds access_token when cross-origin with authToken", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com",
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.formatNavigableHttpURL(
+        "api/files/download",
+        new URLSearchParams({ path: "/tmp/data.csv" }),
+      );
+
+      expect(url.pathname).toBe("/api/files/download");
+      expect(url.searchParams.get("path")).toBe("/tmp/data.csv");
+      expect(url.searchParams.get("access_token")).toBe("my-secret-token");
+    });
+
+    it("omits access_token when same-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: window.location.origin,
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.formatNavigableHttpURL(
+        "api/files/download",
+        new URLSearchParams({ path: "/tmp/data.csv" }),
+      );
+
+      expect(url.searchParams.get("access_token")).toBeNull();
+      expect(url.searchParams.get("path")).toBe("/tmp/data.csv");
+    });
+
+    it("omits access_token when no authToken configured", () => {
+      const runtime = new RuntimeManager(
+        { url: "https://sandbox.example.com", lazy: true },
+        true,
+      );
+      const url = runtime.formatNavigableHttpURL("api/files/download");
+
+      expect(url.searchParams.get("access_token")).toBeNull();
+    });
+
+    it("does not forward a page-inherited access_token when same-origin", () => {
+      const runtime = new RuntimeManager(
+        { url: window.location.origin, lazy: true },
+        true,
+      );
+      const original = window.location.href;
+      window.history.replaceState({}, "", "/?access_token=page-token");
+      try {
+        const url = runtime.formatNavigableHttpURL("api/files/download");
+        expect(url.searchParams.get("access_token")).toBeNull();
+      } finally {
+        window.history.replaceState({}, "", original);
+      }
+    });
+  });
+
   describe("getSseURL", () => {
     it("should return an http(s) URL with the session ID", () => {
       const runtime = new RuntimeManager(mockConfig);

package/src/core/runtime/runtime.ts

@@ -86,6 +86,35 @@ export class RuntimeManager {
     return baseUrl;
   }
 
+  /**
+   * Channels that cannot send custom headers (WebSockets, and browser
+   * navigations such as file downloads) carry the auth token as a query
+   * param instead, but only when cross-origin. Same-origin requests
+   * authenticate via the session cookie, which keeps the token out of
+   * browser history and logs; cross-origin cookies are blocked by browsers,
+   * so the access_token query param is the only option there.
+   */
+  private appendCrossOriginAuth(url: URL): URL {
+    if (!this.isSameOrigin && this.config.authToken) {
+      url.searchParams.set(KnownQueryParams.accessToken, this.config.authToken);
+    }
+    return url;
+  }
+
+  /**
+   * An HTTP URL for requests made by browser navigation (e.g. anchor-click
+   * downloads), which cannot attach auth headers.
+   */
+  formatNavigableHttpURL(path: string, searchParams?: URLSearchParams): URL {
+    const url = this.formatHttpURL({ path, searchParams });
+    // Drop any token inherited from the current page's query params; re-add it
+    // only when cross-origin. Same-origin downloads authenticate via the
+    // session cookie, so a token in the URL would needlessly leak into
+    // proxy/server logs.
+    url.searchParams.delete(KnownQueryParams.accessToken);
+    return this.appendCrossOriginAuth(url);
+  }
+
   formatWsURL(path: string, searchParams?: URLSearchParams): URL {
     // We don't restrict to known query parameters, since mo.query_params()
     // can accept arbitrary parameters.
@@ -94,16 +123,7 @@ export class RuntimeManager {
       searchParams,
       restrictToKnownQueryParams: false,
     });
-
-    // For cross-origin runtimes, pass the auth token as a query parameter.
-    // WebSocket connections cannot send custom headers (no Authorization
-    // header), and cross-origin cookies are blocked by browsers, so the
-    // access_token query param is the only way to authenticate.
-    if (!this.isSameOrigin && this.config.authToken) {
-      url.searchParams.set(KnownQueryParams.accessToken, this.config.authToken);
-    }
-
-    return asWsUrl(url.toString());
+    return asWsUrl(this.appendCrossOriginAuth(url).toString());
   }
 
   /**