@marimo-team/islands 0.25.3-dev5 → 0.25.3-dev51

package/src/plugins/impl/anywidget/__tests__/widget-binding.test.ts

@@ -7,9 +7,12 @@ import { userConfigAtom } from "@/core/config/config";
 import { parseUserConfig } from "@/core/config/config-schema";
 import { initialModeAtom } from "@/core/mode";
 import { store } from "@/core/state/jotai";
+import type { DataURLString } from "@/utils/json/base64";
 import { Model } from "../model";
 import type { Host } from "../host";
-import type { ModelState } from "../types";
+import type { ModelState, WidgetModelId } from "../types";
+import { WidgetRuntime } from "../runtime";
+import { Deferred } from "@/utils/Deferred";
 import { visibleForTesting } from "../widget-binding";
 
 const { WidgetDefRegistry, WidgetBinding } = visibleForTesting;
@@ -30,13 +33,14 @@ const createTestHost = (): Host => ({
 function createBinding<T extends ModelState>(
   widgetDef: AnyWidget<T>,
   model: Model<T>,
-  options: { controller?: AbortController } = {},
+  options: { controller?: AbortController; onDestroy?: () => void } = {},
 ) {
   return WidgetBinding.create({
     widgetDef,
     model,
     createHost: createTestHost,
     controller: options.controller,
+    onDestroy: options.onDestroy,
   });
 }
 
@@ -123,6 +127,25 @@ describe("WidgetDefRegistry", () => {
     promise2.catch(() => undefined);
   });
 
+  it("does not evict a replacement when an invalidated import fails", async () => {
+    const oldImport = getModule(
+      registry,
+      "http://localhost/old.js",
+      "replaced",
+    );
+    registry.invalidate("replaced");
+    const newImport = registry.getModule({
+      jsUrl: "data:text/javascript,export default {}",
+      jsHash: "replaced",
+      kernelAuthored: true,
+    });
+    await expect(oldImport).rejects.toThrow(/untrusted/i);
+    await newImport;
+    expect(getModule(registry, "http://localhost/new.js", "replaced")).toBe(
+      newImport,
+    );
+  });
+
   describe("URL validation", () => {
     it.each([
       // Attack vector: raw <marimo-anywidget data-js-url=...> in markdown
@@ -148,9 +171,225 @@ describe("WidgetDefRegistry", () => {
       ).rejects.not.toThrow(/untrusted/i);
     });
   });
+
+  describe("static module URLs", () => {
+    afterEach(() => {
+      delete window.__MARIMO_STATIC__;
+      vi.restoreAllMocks();
+    });
+
+    it("releases the generated URL when a static module fails to import", async () => {
+      window.__MARIMO_STATIC__ = {
+        files: {
+          "/@file/widget.js":
+            "data:text/javascript;base64,ZXhwb3J0IHt9" as DataURLString,
+        },
+      };
+      const create = vi
+        .spyOn(URL, "createObjectURL")
+        .mockReturnValue("blob:widget-module");
+      const revoke = vi
+        .spyOn(URL, "revokeObjectURL")
+        .mockImplementation(() => undefined);
+
+      await expect(
+        getModule(registry, "./@file/widget.js", "static-module"),
+      ).rejects.toThrow(/blob:/);
+
+      expect(create).toHaveBeenCalledOnce();
+      expect(revoke).toHaveBeenCalledExactlyOnceWith("blob:widget-module");
+    });
+
+    it("does not revoke active modules when their cache entry is invalidated", async () => {
+      window.__MARIMO_STATIC__ = {
+        files: {
+          "/@file/widget.js":
+            "data:text/javascript;base64,ZXhwb3J0IHt9" as DataURLString,
+        },
+      };
+      const oldUrl = "data:text/javascript,export const version = 1";
+      const newUrl = "data:text/javascript,export const version = 2";
+      vi.spyOn(URL, "createObjectURL")
+        .mockReturnValueOnce(oldUrl)
+        .mockReturnValueOnce(newUrl);
+      const revoke = vi
+        .spyOn(URL, "revokeObjectURL")
+        .mockImplementation(() => undefined);
+      const oldLifetime = new AbortController();
+      const newLifetime = new AbortController();
+      const options = { jsUrl: "./@file/widget.js", jsHash: "invalidated" };
+      const oldImport = registry.getModule({
+        ...options,
+        lifetime: oldLifetime.signal,
+      });
+      registry.invalidate(options.jsHash);
+      const newImport = registry.getModule({
+        ...options,
+        lifetime: newLifetime.signal,
+      });
+      await Promise.all([oldImport, newImport]);
+      expect(revoke).not.toHaveBeenCalled();
+      oldLifetime.abort();
+      await vi.waitFor(() =>
+        expect(revoke).toHaveBeenCalledExactlyOnceWith(oldUrl),
+      );
+      const siblingLifetime = new AbortController();
+      expect(
+        registry.getModule({ ...options, lifetime: siblingLifetime.signal }),
+      ).toBe(newImport);
+      newLifetime.abort();
+      siblingLifetime.abort();
+      await vi.waitFor(() => expect(revoke).toHaveBeenLastCalledWith(newUrl));
+    });
+
+    it("waits for a canceled import to settle before revoking its URL", async () => {
+      window.__MARIMO_STATIC__ = {
+        files: {
+          "/@file/widget.js":
+            "data:text/javascript;base64,ZXhwb3J0IHt9" as DataURLString,
+        },
+      };
+      const moduleUrl = "data:text/javascript,export const pending = true";
+      vi.spyOn(URL, "createObjectURL").mockReturnValue(moduleUrl);
+      const revoke = vi
+        .spyOn(URL, "revokeObjectURL")
+        .mockImplementation(() => undefined);
+      const lifetime = new AbortController();
+      const pending = registry.getModule({
+        jsUrl: "./@file/widget.js",
+        jsHash: "pending",
+        lifetime: lifetime.signal,
+      });
+      lifetime.abort();
+      expect(revoke).not.toHaveBeenCalled();
+      await pending;
+      expect(revoke).toHaveBeenCalledExactlyOnceWith(moduleUrl);
+    });
+
+    it("does not revoke an unchanged URL for a missing static file", async () => {
+      window.__MARIMO_STATIC__ = { files: {} };
+      const revoke = vi
+        .spyOn(URL, "revokeObjectURL")
+        .mockImplementation(() => undefined);
+
+      await expect(
+        getModule(registry, "./@file/missing.js", "missing-module"),
+      ).rejects.toThrow(/http:/);
+
+      expect(revoke).not.toHaveBeenCalled();
+    });
+
+    it("keeps a shared module URL until the last widget generation is destroyed", async () => {
+      window.__MARIMO_STATIC__ = {
+        files: {
+          "/@file/widget.js":
+            "data:text/javascript;base64,ZXhwb3J0IHt9" as DataURLString,
+        },
+      };
+      // Node cannot import blob URLs; a data URL exercises successful loading.
+      const moduleUrl = `data:text/javascript,${encodeURIComponent(
+        "export default { initialize: () => ({ moduleUrl: import.meta.url }) }",
+      )}`;
+      vi.spyOn(URL, "createObjectURL").mockReturnValue(moduleUrl);
+      const revoke = vi
+        .spyOn(URL, "revokeObjectURL")
+        .mockImplementation(() => undefined);
+
+      const createRuntime = (id: string) => {
+        const runtime = new WidgetRuntime(id as WidgetModelId, {
+          timeout: 50,
+          isEditMode: () => true,
+          createHost: createTestHost,
+          onModelTimeout: vi.fn(),
+        });
+        runtime.setModel(new Model({}, createMockComm()));
+        runtime.setSpec({ url: "./@file/widget.js", hash: "static-success" });
+        return runtime;
+      };
+      const first = createRuntime("first");
+      const second = createRuntime("second");
+      const widgets = await Promise.all([
+        first.getWidget(),
+        second.getWidget(),
+      ]);
+      expect(widgets.map((widget) => widget.exports)).toEqual([
+        { moduleUrl },
+        { moduleUrl },
+      ]);
+      expect(URL.createObjectURL).toHaveBeenCalledOnce();
+      expect(revoke).not.toHaveBeenCalled();
+
+      first.dispose();
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(revoke).not.toHaveBeenCalled();
+
+      second.dispose();
+      await vi.waitFor(() => {
+        expect(revoke).toHaveBeenCalledExactlyOnceWith(moduleUrl);
+      });
+
+      // A later generation reuses the evaluated module rather than
+      // re-running its top-level side effects, and does not revoke again.
+      const third = createRuntime("third");
+      const widget = await third.getWidget();
+      expect(widget.exports).toEqual({ moduleUrl });
+      expect(URL.createObjectURL).toHaveBeenCalledOnce();
+      third.dispose();
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(revoke).toHaveBeenCalledOnce();
+    });
+  });
 });
 
 describe("WidgetBinding", () => {
+  it("releases module resources after asynchronous generation cleanup", async () => {
+    const cleanup = new Deferred<void>();
+    const onDestroy = vi.fn();
+    const binding = await createBinding(
+      { initialize: () => () => cleanup.promise },
+      new Model({}, createMockComm()),
+      { onDestroy },
+    );
+    const destroying = binding.destroy();
+    expect(onDestroy).not.toHaveBeenCalled();
+    cleanup.resolve();
+    await destroying;
+    await binding.destroy();
+    expect(onDestroy).toHaveBeenCalledOnce();
+  });
+
+  it("retains module resources until canceled initialization and late cleanup settle", async () => {
+    const initialize = new Deferred<() => Promise<void>>();
+    const cleanup = new Deferred<void>();
+    const controller = new AbortController();
+    const onDestroy = vi.fn();
+    const pending = createBinding(
+      { initialize: () => initialize.promise },
+      new Model({}, createMockComm()),
+      { controller, onDestroy },
+    );
+    controller.abort();
+    await expect(pending).rejects.toThrow(/binding destroyed/);
+    expect(onDestroy).not.toHaveBeenCalled();
+    initialize.resolve(() => cleanup.promise);
+    await Promise.resolve();
+    expect(onDestroy).not.toHaveBeenCalled();
+    cleanup.resolve();
+    await vi.waitFor(() => expect(onDestroy).toHaveBeenCalledOnce());
+  });
+
+  it("releases module resources when initialization fails", async () => {
+    const onDestroy = vi.fn();
+    await expect(
+      createBinding(
+        { initialize: () => Promise.reject(new Error("initialize failed")) },
+        new Model({}, createMockComm()),
+        { onDestroy },
+      ),
+    ).rejects.toThrow("initialize failed");
+    await vi.waitFor(() => expect(onDestroy).toHaveBeenCalledOnce());
+  });
+
   it("should run initialize exactly once and expose exports", async () => {
     const model = new Model<ModelState>({ count: 0 }, createMockComm());
     const exports = { setValue: vi.fn() };

package/src/plugins/impl/anywidget/runtime.ts

@@ -206,21 +206,33 @@ export class WidgetRuntime {
     controller: AbortController;
   }): Promise<WidgetBinding<ModelState>> {
     const { spec, model, controller } = options;
-    const mod = await WIDGET_DEF_REGISTRY.getModule({
-      jsUrl: spec.url,
-      jsHash: spec.hash,
-      kernelAuthored: true,
-    });
-    const widget = resolveAnyWidget(mod, spec.url);
-    if (!widget) {
-      WIDGET_DEF_REGISTRY.invalidate(spec.hash);
-      throw getInvalidAnyWidgetModuleError(mod, spec.url);
+    const moduleLifetime = new AbortController();
+    let widget: NonNullable<ReturnType<typeof resolveAnyWidget>>;
+    try {
+      const mod = await WIDGET_DEF_REGISTRY.getModule({
+        jsUrl: spec.url,
+        jsHash: spec.hash,
+        kernelAuthored: true,
+        lifetime: moduleLifetime.signal,
+      });
+      const resolved = resolveAnyWidget(mod, spec.url);
+      if (!resolved) {
+        WIDGET_DEF_REGISTRY.invalidate(spec.hash);
+        throw getInvalidAnyWidgetModuleError(mod, spec.url);
+      }
+      widget = resolved;
+    } catch (error) {
+      moduleLifetime.abort();
+      throw error;
     }
+    // From here WidgetBinding owns release, including late cleanup when the
+    // generation is destroyed during initialize.
     return WidgetBinding.create({
       widgetDef: widget,
       model,
       createHost: this.#createHost,
       controller,
+      onDestroy: () => moduleLifetime.abort(),
     });
   }
 

package/src/plugins/impl/anywidget/widget-binding.ts

@@ -46,38 +46,80 @@ function abortSignalAny(signals: AbortSignal[]): AbortSignal {
   return controller.signal;
 }
 
+interface ModuleEntry {
+  promise: Promise<unknown>;
+  references: number;
+  dispose: () => void;
+}
+
 /**
  * Deduplicates ESM imports by jsHash.
  * A single import is shared across all widget instances using the same module.
  */
 class WidgetDefRegistry {
-  #cache = new Map<string, Promise<unknown>>();
+  #cache = new Map<string, ModuleEntry>();
 
   /**
    * Get (or start) the ESM import for a widget module, cached by
    * jsHash so multiple instances share one import. Pass
    * `kernelAuthored: true` only for URLs from an `EsmSpec`; it widens
-   * the import gate to remote and data URLs.
+   * the import gate to remote and data URLs. A lifetime retains the module's
+   * temporary URL until that signal aborts; omit it to retain the URL for
+   * the notebook's lifetime. The evaluated module itself stays cached.
    */
   public getModule(options: {
     jsUrl: string;
     jsHash: string;
     kernelAuthored?: boolean;
+    lifetime?: AbortSignal;
   }): Promise<unknown> {
-    const { jsUrl, jsHash, kernelAuthored = false } = options;
-    const cached = this.#cache.get(jsHash);
-    if (cached) {
-      return cached;
+    const { jsUrl, jsHash, kernelAuthored = false, lifetime } = options;
+    let entry = this.#cache.get(jsHash);
+    if (!entry) {
+      const imported = this.#doImport(jsUrl, { kernelAuthored });
+      const created: ModuleEntry = {
+        references: 0,
+        dispose: () => undefined,
+        promise: imported.then(
+          ({ module, dispose }) => {
+            created.dispose = dispose;
+            return module;
+          },
+          (error) => {
+            // An invalidated import must not evict its replacement.
+            if (this.#cache.get(jsHash) === created) {
+              this.#cache.delete(jsHash);
+            }
+            throw error;
+          },
+        ),
+      };
+      this.#cache.set(jsHash, created);
+      entry = created;
     }
-
-    const promise = this.#doImport(jsUrl, { kernelAuthored }).catch((error) => {
-      // On failure, remove from cache so a retry with a new URL can work
-      this.#cache.delete(jsHash);
-      throw error;
-    });
-
-    this.#cache.set(jsHash, promise);
-    return promise;
+    const retained = entry;
+    retained.references += 1;
+    const release = () => {
+      retained.references -= 1;
+      if (retained.references === 0) {
+        // Keep the evaluated module cached: re-importing would re-run its
+        // top-level side effects (e.g. `customElements.define`). Only the
+        // temporary URL is released.
+        void retained.promise.then(
+          () => {
+            retained.dispose();
+            retained.dispose = () => undefined;
+          },
+          () => undefined,
+        );
+      }
+    };
+    if (lifetime?.aborted) {
+      release();
+    } else {
+      lifetime?.addEventListener("abort", release, { once: true });
+    }
+    return retained.promise;
   }
 
   public invalidate(jsHash: string): void {
@@ -87,7 +129,7 @@ class WidgetDefRegistry {
   async #doImport(
     jsUrl: string,
     opts: { kernelAuthored?: boolean } = {},
-  ): Promise<unknown> {
+  ): Promise<{ module: unknown; dispose: () => void }> {
     // By default, only trust marimo virtual file paths: arbitrary URLs
     // would let markdown-injected elements import attacker-controlled
     // JavaScript. Kernel-authored URLs (from an `EsmSpec`) are exempt —
@@ -103,11 +145,23 @@ class WidgetDefRegistry {
         )}`,
       );
     }
-    let url = asRemoteURL(jsUrl).toString();
-    if (isStaticNotebook()) {
-      url = resolveVirtualFileURL(url);
+    const remoteUrl = asRemoteURL(jsUrl).toString();
+    const url = isStaticNotebook()
+      ? resolveVirtualFileURL(remoteUrl)
+      : remoteUrl;
+    const dispose = () => {
+      // Only the static-file resolver's newly created URL belongs to us.
+      if (url !== remoteUrl) {
+        URL.revokeObjectURL(url);
+      }
+    };
+    try {
+      const module = await import(/* @vite-ignore */ url);
+      return { module, dispose };
+    } catch (error) {
+      dispose();
+      throw error;
     }
-    return import(/* @vite-ignore */ url);
   }
 }
 
@@ -132,6 +186,7 @@ interface CreateWidgetBindingOptions<T extends ModelState> {
   model: Model<T>;
   createHost: HostFactory;
   controller?: AbortController;
+  onDestroy?: () => void;
 }
 
 async function safelyRunCleanup(cleanup: Cleanup, reason: string) {
@@ -159,6 +214,7 @@ export class WidgetBinding<T extends ModelState = ModelState> {
   #createHost: HostFactory;
   #cleanupTasks = new Set<Promise<void>>();
   #viewTasks = new Set<Promise<void>>();
+  #onDestroy?: () => void;
 
   private constructor(options: {
     widget: ResolvedWidget<T>;
@@ -166,12 +222,14 @@ export class WidgetBinding<T extends ModelState = ModelState> {
     exports: unknown;
     controller: AbortController;
     createHost: HostFactory;
+    onDestroy?: () => void;
   }) {
     this.#widget = options.widget;
     this.#model = options.model;
     this.#exports = options.exports;
     this.#controller = options.controller;
     this.#createHost = options.createHost;
+    this.#onDestroy = options.onDestroy;
   }
 
   /**
@@ -185,76 +243,94 @@ export class WidgetBinding<T extends ModelState = ModelState> {
     model,
     createHost,
     controller = new AbortController(),
+    onDestroy,
   }: CreateWidgetBindingOptions<T>): Promise<WidgetBinding<T>> {
     const signal = controller.signal;
-    const widget: ResolvedWidget<T> =
-      typeof widgetDef === "function" ? await widgetDef() : widgetDef;
+    let lateCleanup: Promise<void> | undefined;
+    try {
+      const widget: ResolvedWidget<T> =
+        typeof widgetDef === "function" ? await widgetDef() : widgetDef;
 
-    const initPromise = Promise.resolve(
-      widget.initialize?.({
-        model: modelProxy(model, signal),
-        experimental,
-        signal,
-      }),
-    );
-    // If destroyed mid-initialize, still run a late-arriving legacy
-    // cleanup callback: the widget acquired resources.
-    initPromise
-      .then(async (settled) => {
-        if (signal.aborted && isCleanup(settled)) {
-          await safelyRunCleanup(settled, "late initialize");
-        }
-      })
-      .catch(() => undefined);
+      if (signal.aborted) {
+        throw new Error("[anywidget] binding destroyed");
+      }
+      const initPromise = Promise.resolve(
+        widget.initialize?.({
+          model: modelProxy(model, signal),
+          experimental,
+          signal,
+        }),
+      );
+      // If destroyed mid-initialize, still run a late-arriving legacy
+      // cleanup callback: the widget acquired resources.
+      lateCleanup = initPromise
+        .then(async (settled) => {
+          if (signal.aborted && isCleanup(settled)) {
+            await safelyRunCleanup(settled, "late initialize");
+          }
+        })
+        .catch(() => undefined);
 
-    // Race against destruction so a hung `initialize` can't strand
-    // callers whose model already closed.
-    const result = await Promise.race([
-      initPromise,
-      new Promise<never>((_, reject) => {
-        const abort = () => reject(new Error("[anywidget] binding destroyed"));
-        if (signal.aborted) {
-          abort();
-          return;
-        }
-        signal.addEventListener("abort", abort, { once: true });
-      }),
-    ]);
+      // Race against destruction so a hung `initialize` can't strand
+      // callers whose model already closed.
+      const result = await Promise.race([
+        initPromise,
+        new Promise<never>((_, reject) => {
+          const abort = () =>
+            reject(new Error("[anywidget] binding destroyed"));
+          if (signal.aborted) {
+            abort();
+            return;
+          }
+          signal.addEventListener("abort", abort, { once: true });
+        }),
+      ]);
 
-    if (signal.aborted) {
-      throw new Error("[anywidget] binding destroyed");
-    }
+      if (signal.aborted) {
+        throw new Error("[anywidget] binding destroyed");
+      }
 
-    // Distinguish anywidget's three return shapes:
-    //   function → legacy cleanup callback (run on destroy)
-    //   object   → widget exports (anywidget>=0.11)
-    //   void     → nothing
-    let exports: unknown;
-    let initializeCleanup: Cleanup | undefined;
-    if (isCleanup(result)) {
-      initializeCleanup = result;
-      exports = undefined;
-    } else if (typeof result === "object" && result !== null) {
-      exports = result;
-    } else {
-      exports = undefined;
-    }
+      // Distinguish anywidget's three return shapes:
+      //   function → legacy cleanup callback (run on destroy)
+      //   object   → widget exports (anywidget>=0.11)
+      //   void     → nothing
+      let exports: unknown;
+      let initializeCleanup: Cleanup | undefined;
+      if (isCleanup(result)) {
+        initializeCleanup = result;
+        exports = undefined;
+      } else if (typeof result === "object" && result !== null) {
+        exports = result;
+      } else {
+        exports = undefined;
+      }
 
-    const binding = new WidgetBinding({
-      widget,
-      model,
-      exports,
-      controller,
-      createHost,
-    });
-    if (initializeCleanup) {
-      signal.addEventListener(
-        "abort",
-        () => void binding.#trackCleanup(initializeCleanup, "initialize"),
-        { once: true },
-      );
+      const binding = new WidgetBinding({
+        widget,
+        model,
+        exports,
+        controller,
+        createHost,
+        onDestroy,
+      });
+      if (initializeCleanup) {
+        signal.addEventListener(
+          "abort",
+          () => void binding.#trackCleanup(initializeCleanup, "initialize"),
+          { once: true },
+        );
+      }
+      return binding;
+    } catch (error) {
+      // A canceled initialize may still be using import.meta.url in its
+      // late cleanup. Keep the module alive until that work settles.
+      if (lateCleanup) {
+        void lateCleanup.then(() => onDestroy?.());
+      } else {
+        onDestroy?.();
+      }
+      throw error;
     }
-    return binding;
   }
 
   /**
@@ -353,6 +429,8 @@ export class WidgetBinding<T extends ModelState = ModelState> {
     while (this.#cleanupTasks.size > 0) {
       await Promise.all(this.#cleanupTasks);
     }
+    this.#onDestroy?.();
+    this.#onDestroy = undefined;
   }
 }
 

package/src/plugins/impl/code/any-language-editor.tsx

@@ -53,14 +53,14 @@ const AnyLanguageCodeMirror: React.FC<
     extensions?: readonly Extension[];
   }
 > = ({
-  language,
+  language: requestedLanguage,
   hideUnsupportedLanguageErrors,
   showCopyButton,
   extensions = EMPTY_EXTENSIONS,
   ...props
 }) => {
   // Maybe normalize the language to the extension
-  language = LANGUAGE_MAP[language || ""] || language;
+  const language = LANGUAGE_MAP[requestedLanguage || ""] || requestedLanguage;
 
   const isSupported = isSupportedLanguage(language);
   if (!isSupported) {