@marimo-team/frontend 0.19.9-dev7 → 0.19.9-dev9

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

@@ -0,0 +1,240 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { Model } from "../model";
+import type { ModelState, WidgetModelId } from "../types";
+import { visibleForTesting } from "../widget-binding";
+
+const { WidgetDefRegistry, WidgetBinding, BindingManager } = visibleForTesting;
+
+// Helper to create typed model IDs for tests
+const asModelId = (id: string): WidgetModelId => id as WidgetModelId;
+
+function createMockComm() {
+  return {
+    sendUpdate: vi.fn().mockResolvedValue(undefined),
+    sendCustomMessage: vi.fn().mockResolvedValue(undefined),
+  };
+}
+
+describe("WidgetDefRegistry", () => {
+  let registry: InstanceType<typeof WidgetDefRegistry>;
+
+  beforeEach(() => {
+    registry = new WidgetDefRegistry();
+  });
+
+  it("should cache modules by jsHash and return same promise", () => {
+    // Two calls with same hash should return the exact same promise object
+    const promise1 = registry.getModule("http://localhost/widget.js", "hash1");
+    const promise2 = registry.getModule("http://localhost/widget.js", "hash1");
+    expect(promise1).toBe(promise2);
+    // Catch the unhandled rejection from the import() attempt
+    promise1.catch(() => undefined);
+  });
+
+  it("should deduplicate concurrent imports for the same hash", () => {
+    const promise1 = registry.getModule("http://localhost/a.js", "same-hash");
+    const promise2 = registry.getModule("http://localhost/b.js", "same-hash");
+    // Same hash means same promise, even with different URLs
+    expect(promise1).toBe(promise2);
+    promise1.catch(() => undefined);
+  });
+
+  it("should create different promises for different hashes", () => {
+    const promise1 = registry.getModule("http://localhost/a.js", "hash-a");
+    const promise2 = registry.getModule("http://localhost/b.js", "hash-b");
+    expect(promise1).not.toBe(promise2);
+    promise1.catch(() => undefined);
+    promise2.catch(() => undefined);
+  });
+
+  it("should invalidate cached modules", () => {
+    const promise1 = registry.getModule("http://localhost/a.js", "hash1");
+    promise1.catch(() => undefined);
+    registry.invalidate("hash1");
+    const promise2 = registry.getModule("http://localhost/a.js", "hash1");
+    promise2.catch(() => undefined);
+    expect(promise1).not.toBe(promise2);
+  });
+
+  it("should remove from cache on import failure so retry creates new promise", async () => {
+    const promise1 = registry.getModule("http://localhost/a.js", "fail-hash");
+    // The import will fail in Node (http: scheme not supported)
+    await expect(promise1).rejects.toThrow();
+    // After failure, cache should be cleared, so next call creates a new promise
+    const promise2 = registry.getModule("http://localhost/a.js", "fail-hash");
+    expect(promise1).not.toBe(promise2);
+    promise2.catch(() => undefined);
+  });
+});
+
+describe("WidgetBinding", () => {
+  let binding: InstanceType<typeof WidgetBinding>;
+  let model: Model<ModelState>;
+
+  beforeEach(() => {
+    binding = new WidgetBinding();
+    model = new Model<ModelState>({ count: 0 }, createMockComm());
+  });
+
+  it("should initialize once and return a render function", async () => {
+    const initCleanup = vi.fn();
+    const renderCleanup = vi.fn();
+    const widgetDef = {
+      initialize: vi.fn().mockResolvedValue(initCleanup),
+      render: vi.fn().mockResolvedValue(renderCleanup),
+    };
+
+    const renderFn = await binding.bind(widgetDef, model);
+    expect(widgetDef.initialize).toHaveBeenCalledTimes(1);
+    expect(typeof renderFn).toBe("function");
+
+    // Render into an element
+    const el = document.createElement("div");
+    const controller = new AbortController();
+    await renderFn(el, controller.signal);
+    expect(widgetDef.render).toHaveBeenCalledTimes(1);
+  });
+
+  it("should return cached render for same widget def", async () => {
+    const widgetDef = {
+      initialize: vi.fn(),
+      render: vi.fn(),
+    };
+
+    const render1 = await binding.bind(widgetDef, model);
+    const render2 = await binding.bind(widgetDef, model);
+    expect(render1).toBe(render2);
+    // Initialize should only be called once
+    expect(widgetDef.initialize).toHaveBeenCalledTimes(1);
+  });
+
+  it("should re-initialize on hot reload (different widget def)", async () => {
+    const cleanup1 = vi.fn();
+    const widgetDef1 = {
+      initialize: vi.fn().mockResolvedValue(cleanup1),
+      render: vi.fn(),
+    };
+
+    const widgetDef2 = {
+      initialize: vi.fn(),
+      render: vi.fn(),
+    };
+
+    const render1 = await binding.bind(widgetDef1, model);
+    const render2 = await binding.bind(widgetDef2, model);
+
+    expect(render1).not.toBe(render2);
+    expect(cleanup1).toHaveBeenCalledTimes(1); // Old binding cleaned up
+    expect(widgetDef2.initialize).toHaveBeenCalledTimes(1);
+  });
+
+  it("should cleanup render on view signal abort", async () => {
+    const renderCleanup = vi.fn();
+    const widgetDef = {
+      initialize: vi.fn(),
+      render: vi.fn().mockResolvedValue(renderCleanup),
+    };
+
+    const renderFn = await binding.bind(widgetDef, model);
+    const el = document.createElement("div");
+    const viewController = new AbortController();
+    await renderFn(el, viewController.signal);
+
+    // Aborting the view signal should trigger render cleanup
+    viewController.abort();
+    expect(renderCleanup).toHaveBeenCalledTimes(1);
+  });
+
+  it("should cleanup everything on destroy", async () => {
+    const initCleanup = vi.fn();
+    const renderCleanup = vi.fn();
+    const widgetDef = {
+      initialize: vi.fn().mockResolvedValue(initCleanup),
+      render: vi.fn().mockResolvedValue(renderCleanup),
+    };
+
+    const renderFn = await binding.bind(widgetDef, model);
+    const el = document.createElement("div");
+    const viewController = new AbortController();
+    await renderFn(el, viewController.signal);
+
+    binding.destroy();
+    expect(initCleanup).toHaveBeenCalledTimes(1);
+    expect(renderCleanup).toHaveBeenCalledTimes(1);
+  });
+
+  it("should handle widget def as a function", async () => {
+    const widget = {
+      initialize: vi.fn(),
+      render: vi.fn(),
+    };
+    const widgetDefFn = vi.fn().mockResolvedValue(widget);
+
+    await binding.bind(widgetDefFn, model);
+    expect(widgetDefFn).toHaveBeenCalledTimes(1);
+    expect(widget.initialize).toHaveBeenCalledTimes(1);
+  });
+
+  it("should handle widget with no initialize or render", async () => {
+    const widgetDef = {};
+    const renderFn = await binding.bind(widgetDef, model);
+    expect(typeof renderFn).toBe("function");
+
+    // Render should not throw
+    const el = document.createElement("div");
+    const controller = new AbortController();
+    await renderFn(el, controller.signal);
+  });
+});
+
+describe("BindingManager", () => {
+  let manager: InstanceType<typeof BindingManager>;
+
+  beforeEach(() => {
+    manager = new BindingManager();
+  });
+
+  it("should create bindings on demand", () => {
+    const modelId = asModelId("model-1");
+    expect(manager.has(modelId)).toBe(false);
+
+    const binding = manager.getOrCreate(modelId);
+    expect(binding).toBeDefined();
+    expect(manager.has(modelId)).toBe(true);
+  });
+
+  it("should return the same binding for the same model id", () => {
+    const modelId = asModelId("model-1");
+    const binding1 = manager.getOrCreate(modelId);
+    const binding2 = manager.getOrCreate(modelId);
+    expect(binding1).toBe(binding2);
+  });
+
+  it("should destroy and remove bindings", async () => {
+    const modelId = asModelId("model-1");
+    const binding = manager.getOrCreate(modelId);
+
+    const model = new Model<ModelState>({ count: 0 }, createMockComm());
+    const initCleanup = vi.fn();
+    const widgetDef = {
+      initialize: vi.fn().mockResolvedValue(initCleanup),
+      render: vi.fn(),
+    };
+    await binding.bind(widgetDef, model);
+
+    manager.destroy(modelId);
+    expect(manager.has(modelId)).toBe(false);
+    expect(initCleanup).toHaveBeenCalledTimes(1);
+  });
+
+  it("should handle idempotent destroy", () => {
+    const modelId = asModelId("model-1");
+    manager.getOrCreate(modelId);
+
+    // Should not throw
+    manager.destroy(modelId);
+    manager.destroy(modelId);
+    expect(manager.has(modelId)).toBe(false);
+  });
+});

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

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
-import type { AnyModel, AnyWidget, Experimental } from "@anywidget/types";
+import type { AnyModel } from "@anywidget/types";
 import { debounce } from "lodash-es";
 import type { NotificationMessageData } from "@/core/kernel/messages";
 import { getRequestClient } from "@/core/network/requests";
@@ -20,62 +20,62 @@ import { Logger } from "@/utils/Logger";
 import { repl } from "@/utils/repl";
 import type { AnyWidgetMessage } from "./schemas";
 import type { EventHandler, ModelState, WidgetModelId } from "./types";
+import { BINDING_MANAGER } from "./widget-binding";
+
+interface ModelEntry {
+  deferred: Deferred<Model<ModelState>>;
+  controller: AbortController;
+}
 
 class ModelManager {
-  /**
-   * Map of model ids to deferred promises
-   */
-  #models = new Map<WidgetModelId, Deferred<Model<ModelState>>>();
-  /**
-   * Timeout for model lookup
-   */
+  #entries = new Map<WidgetModelId, ModelEntry>();
   #timeout: number;
 
   constructor(timeout = 10_000) {
     this.#timeout = timeout;
   }
 
-  get(key: WidgetModelId): Promise<Model<any>> {
-    let deferred = this.#models.get(key);
-    if (deferred) {
-      return deferred.promise;
+  #getOrCreateEntry(key: WidgetModelId): ModelEntry {
+    let entry = this.#entries.get(key);
+    if (!entry) {
+      entry = {
+        deferred: new Deferred<Model<ModelState>>(),
+        controller: new AbortController(),
+      };
+      this.#entries.set(key, entry);
     }
-
-    // If the model is not yet created, create the new deferred promise without resolving it
-    deferred = new Deferred<Model<ModelState>>();
-    this.#models.set(key, deferred);
-
-    // Add timeout to prevent hanging
-    setTimeout(() => {
-      // Already settled
-      if (deferred.status !== "pending") {
-        return;
-      }
-
-      deferred.reject(new Error(`Model not found for key: ${key}`));
-      this.#models.delete(key);
-    }, this.#timeout);
-
-    return deferred.promise;
+    return entry;
   }
 
-  set(key: WidgetModelId, model: Model<any>): void {
-    let deferred = this.#models.get(key);
-    if (!deferred) {
-      deferred = new Deferred<Model<ModelState>>();
-      this.#models.set(key, deferred);
+  get(key: WidgetModelId): Promise<Model<any>> {
+    const entry = this.#getOrCreateEntry(key);
+    if (entry.deferred.status === "pending") {
+      // Add timeout to prevent hanging
+      setTimeout(() => {
+        if (entry.deferred.status !== "pending") {
+          return;
+        }
+        entry.deferred.reject(new Error(`Model not found for key: ${key}`));
+        this.#entries.delete(key);
+      }, this.#timeout);
     }
-    deferred.resolve(model);
+    return entry.deferred.promise;
   }
 
   /**
-   * Check if a model exists and has been resolved (not pending).
-   * This is useful for checking if a model was already created by the plugin
-   * before the 'open' message arrives.
+   * Create a model with a managed lifecycle signal.
+   * The signal is aborted when the model is deleted.
    */
-  has(key: WidgetModelId): boolean {
-    const deferred = this.#models.get(key);
-    return deferred !== undefined && deferred.status === "resolved";
+  create(
+    key: WidgetModelId,
+    factory: (signal: AbortSignal) => Model<ModelState>,
+  ): void {
+    const entry = this.#getOrCreateEntry(key);
+    entry.deferred.resolve(factory(entry.controller.signal));
+  }
+
+  set(key: WidgetModelId, model: Model<any>): void {
+    this.#getOrCreateEntry(key).deferred.resolve(model);
   }
 
   /**
@@ -83,15 +83,19 @@ class ModelManager {
    * Returns undefined if the model doesn't exist or is still pending.
    */
   getSync(key: WidgetModelId): Model<any> | undefined {
-    const deferred = this.#models.get(key);
-    if (deferred && deferred.status === "resolved") {
-      return deferred.value;
+    const entry = this.#entries.get(key);
+    if (entry && entry.deferred.status === "resolved") {
+      return entry.deferred.value;
     }
     return undefined;
   }
 
   delete(key: WidgetModelId): void {
-    this.#models.delete(key);
+    Logger.debug(
+      `[ModelManager] Deleting model=${key}, aborting lifecycle signal`,
+    );
+    this.#entries.get(key)?.controller.abort();
+    this.#entries.delete(key);
   }
 }
 
@@ -102,28 +106,7 @@ interface MarimoComm<T> {
 
 const marimoSymbol = Symbol("marimo");
 
-const experimental: Experimental = {
-  invoke: async () => {
-    const message =
-      "anywidget.invoke not supported in marimo. Please file an issue at https://github.com/marimo-team/marimo/issues";
-    Logger.warn(message);
-    throw new Error(message);
-  },
-};
-
-type RenderFn = (el: HTMLElement, signal: AbortSignal) => Promise<void>;
-
 interface MarimoInternalApi<T extends ModelState> {
-  /**
-   * Resolve the widget definition and initialize if needed.
-   * Returns a render function that can be called for each view.
-   *
-   * Per AFM spec:
-   * - widgetDef() is called once per model
-   * - initialize() is called once per model
-   * - render() (the returned function) is called once per view
-   */
-  resolveWidget: (widgetDef: AnyWidget<T>) => Promise<RenderFn>;
   /**
    * Update model state and emit change events for any differences.
    */
@@ -135,10 +118,6 @@ interface MarimoInternalApi<T extends ModelState> {
     message: Extract<AnyWidgetMessage, { method: "custom" }>,
     buffers?: readonly DataView[],
   ) => void;
-  /**
-   * Destroy the model, triggering initialize cleanup.
-   */
-  destroy: () => void;
 }
 
 /**
@@ -159,18 +138,19 @@ export class Model<T extends ModelState> implements AnyModel<T> {
   #data: T;
   #comm: MarimoComm<T>;
   #listeners: Record<string, Set<EventHandler> | undefined> = {};
-  #controller = new AbortController();
-  #widgetDef: AnyWidget<T> | undefined;
-  #render:
-    | ((el: HTMLElement, signal: AbortSignal) => Promise<void>)
-    | undefined;
 
   static _modelManager: ModelManager = MODEL_MANAGER;
 
-  constructor(data: T, comm: MarimoComm<T>) {
+  constructor(data: T, comm: MarimoComm<T>, signal?: AbortSignal) {
     this.#data = data;
     this.#comm = comm;
     this.#dirtyFields = new Map();
+    if (signal) {
+      signal.addEventListener("abort", () => {
+        Logger.debug("[Model] Signal aborted, clearing all listeners");
+        this.#listeners = {};
+      });
+    }
   }
 
   /**
@@ -183,52 +163,6 @@ export class Model<T extends ModelState> implements AnyModel<T> {
       message: Extract<AnyWidgetMessage, { method: "custom" }>,
       buffers?: readonly DataView[],
     ) => this.#emitCustomMessage(message, buffers),
-    resolveWidget: async (widgetDef: AnyWidget<T>): Promise<RenderFn> => {
-      // Already initialized with the same widget - return cached render
-      if (this.#render && this.#widgetDef === widgetDef) {
-        return this.#render;
-      }
-
-      // If widgetDef changed (hot reload), destroy old and re-initialize
-      if (this.#render && this.#widgetDef !== widgetDef) {
-        this.#controller.abort();
-        this.#controller = new AbortController();
-        this.#render = undefined;
-      }
-
-      this.#widgetDef = widgetDef;
-
-      // Resolve the widget definition (call if it's a function)
-      const widget =
-        typeof widgetDef === "function" ? await widgetDef() : widgetDef;
-
-      // Call initialize once per model
-      const cleanup = await widget.initialize?.({ model: this, experimental });
-      if (cleanup) {
-        this.#controller.signal.addEventListener("abort", cleanup);
-      }
-
-      // Store and return the render closure
-      this.#render = async (el: HTMLElement, signal: AbortSignal) => {
-        const renderCleanup = await widget.render?.({
-          model: this,
-          el,
-          experimental,
-        });
-        if (renderCleanup) {
-          // Cleanup when either the view unmounts or the model is destroyed
-          AbortSignal.any([signal, this.#controller.signal]).addEventListener(
-            "abort",
-            renderCleanup,
-          );
-        }
-      };
-
-      return this.#render;
-    },
-    destroy: () => {
-      this.#controller.abort();
-    },
   };
 
   off(eventName?: string | null, callback?: EventHandler | null): void {
@@ -387,8 +321,6 @@ export async function handleWidgetMessage(
   const modelId = notification.model_id as WidgetModelId;
   const msg = notification.message;
 
-  Logger.debug("AnyWidget message", msg);
-
   // Decode base64 buffers to DataViews (present in open/update/custom messages)
   const base64Buffers: Base64String[] =
     "buffers" in msg ? (msg.buffers as Base64String[]) : [];
@@ -411,25 +343,44 @@ export async function handleWidgetMessage(
         return;
       }
 
-      const model = new Model(stateWithBuffers, {
-        async sendUpdate(changeData) {
-          const { state, buffers, bufferPaths } =
-            serializeBuffersToBase64(changeData);
-          await getRequestClient().sendModelValue({
-            modelId,
-            message: { method: "update", state, bufferPaths },
-            buffers,
-          });
-        },
-        async sendCustomMessage(content, buffers) {
-          await getRequestClient().sendModelValue({
-            modelId,
-            message: { method: "custom", content },
-            buffers: buffers.map(dataViewToBase64),
-          });
-        },
-      });
-      modelManager.set(modelId, model);
+      modelManager.create(
+        modelId,
+        (signal) =>
+          new Model(
+            stateWithBuffers,
+            {
+              async sendUpdate(changeData) {
+                if (signal.aborted) {
+                  Logger.debug(
+                    `[Model] sendUpdate suppressed for model=${modelId} (signal aborted)`,
+                  );
+                  return;
+                }
+                const { state, buffers, bufferPaths } =
+                  serializeBuffersToBase64(changeData);
+                await getRequestClient().sendModelValue({
+                  modelId,
+                  message: { method: "update", state, bufferPaths },
+                  buffers,
+                });
+              },
+              async sendCustomMessage(content, buffers) {
+                if (signal.aborted) {
+                  Logger.debug(
+                    `[Model] sendCustomMessage suppressed for model=${modelId} (signal aborted)`,
+                  );
+                  return;
+                }
+                await getRequestClient().sendModelValue({
+                  modelId,
+                  message: { method: "custom", content },
+                  buffers: buffers.map(dataViewToBase64),
+                });
+              },
+            },
+            signal,
+          ),
+      );
       return;
     }
 
@@ -444,11 +395,8 @@ export async function handleWidgetMessage(
     }
 
     case "close": {
-      const model = modelManager.getSync(modelId);
-      if (model) {
-        getMarimoInternal(model).destroy();
-      }
-      modelManager.delete(modelId);
+      BINDING_MANAGER.destroy(modelId);
+      modelManager.delete(modelId); // aborts the model's signal, clearing listeners
       return;
     }