@marimo-team/islands 0.23.14-dev6 → 0.23.14-dev61

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

@@ -0,0 +1,70 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { AnyModel } from "@anywidget/types";
+import type { ModelState } from "./types";
+
+type ModelEventCallback = Parameters<AnyModel<ModelState>["on"]>[1];
+
+/**
+ * A listener registered through a model proxy (see the hydration
+ * replay in `WidgetBinding.createView`).
+ */
+export interface ProxyRegistration {
+  event: string;
+  callback: ModelEventCallback;
+}
+
+/**
+ * Wrap a model so every `on()` call from inside `initialize` or `render`
+ * is auto-tied to a lifetime `AbortSignal`. When the signal aborts,
+ * every listener registered through the proxy is removed.
+ *
+ * Without this, a parent widget that re-renders a child via
+ * `host.getWidget().render` would accumulate listeners on the child's
+ * model on every re-render — the widget author would have to know to
+ * pass `signal` to each `on()` call themselves.
+ *
+ * The proxy is purely a host-side ergonomic helper. The widget author
+ * still writes `model.on("change:foo", handler)` exactly as before; the
+ * cleanup signal is supplied transparently.
+ *
+ * `onRegister`, when given, observes each `on()` registration.
+ */
+// oxlint-disable-next-line marimo/prefer-object-params -- concise internal helper used at protocol call sites
+export function modelProxy<T extends ModelState>(
+  model: AnyModel<T>,
+  signal: AbortSignal,
+  onRegister?: (registration: ProxyRegistration) => void,
+): AnyModel<T> {
+  return {
+    get(key) {
+      return model.get(key);
+    },
+    set(key, value) {
+      model.set(key, value);
+    },
+    save_changes() {
+      model.save_changes();
+    },
+    send: model.send.bind(model),
+    on(name: string, callback: ModelEventCallback): void {
+      if (signal.aborted) {
+        return;
+      }
+      model.on(name, callback);
+      signal.addEventListener("abort", () => model.off(name, callback), {
+        once: true,
+      });
+      onRegister?.({ event: name, callback });
+    },
+    off(name?: string | null, callback?: ModelEventCallback | null): void {
+      model.off(name ?? null, callback ?? null);
+    },
+    widget_manager: {
+      async get_model<TT extends ModelState>(modelId: string) {
+        const child = await model.widget_manager.get_model<TT>(modelId);
+        return modelProxy(child, signal);
+      },
+    },
+  };
+}

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

@@ -3,104 +3,24 @@
 
 import type { AnyModel } from "@anywidget/types";
 import { debounce } from "lodash-es";
-import type { NotificationMessageData } from "@/core/kernel/messages";
-import { getRequestClient } from "@/core/network/requests";
-import { isStaticNotebook } from "@/core/static/static-state";
-import {
-  decodeFromWire,
-  serializeBuffersToBase64,
-} from "@/plugins/impl/anywidget/serialization";
-import { assertNever } from "@/utils/assertNever";
-import { Deferred } from "@/utils/Deferred";
-import {
-  type Base64String,
-  base64ToDataView,
-  dataViewToBase64,
-} from "@/utils/json/base64";
 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 {
-  #entries = new Map<WidgetModelId, ModelEntry>();
-  #timeout: number;
-
-  constructor(timeout = 10_000) {
-    this.#timeout = timeout;
-  }
-
-  #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);
-    }
-    return entry;
-  }
-
-  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);
-    }
-    return entry.deferred.promise;
-  }
-
-  /**
-   * Create a model with a managed lifecycle signal.
-   * The signal is aborted when the model is deleted.
-   */
-  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);
-  }
-
-  /**
-   * Get a model synchronously if it exists and has been resolved.
-   * Returns undefined if the model doesn't exist or is still pending.
-   */
-  getSync(key: WidgetModelId): Model<any> | undefined {
-    const entry = this.#entries.get(key);
-    if (entry && entry.deferred.status === "resolved") {
-      return entry.deferred.value;
-    }
-    return undefined;
-  }
-
-  delete(key: WidgetModelId): void {
-    Logger.debug(
-      `[ModelManager] Deleting model=${key}, aborting lifecycle signal`,
-    );
-    this.#entries.get(key)?.controller.abort();
-    this.#entries.delete(key);
-  }
+/**
+ * The custom-message shape delivered to `msg:custom` listeners — the
+ * only message shape the frontend re-declares; the wire contract is
+ * the generated `model-lifecycle` type.
+ */
+export interface CustomMessage {
+  method: "custom";
+  // oxlint-disable-next-line typescript/no-explicit-any
+  content: any;
 }
 
-interface MarimoComm<T> {
+/**
+ * The model's channel back to the kernel; a no-op in static exports.
+ */
+export interface MarimoComm<T> {
   sendUpdate: (value: Partial<T>) => Promise<void>;
   sendCustomMessage: (content: unknown, buffers: DataView[]) => Promise<void>;
 }
@@ -112,15 +32,11 @@ interface MarimoInternalApi<T extends ModelState> {
    * Update model state and emit change events for any differences.
    */
   updateAndEmitDiffs: (value: T) => void;
-  /**
-   * Re-emit current state as change events.
-   */
-  reemitState: () => void;
   /**
    * Emit a custom message to listeners.
    */
   emitCustomMessage: (
-    message: Extract<AnyWidgetMessage, { method: "custom" }>,
+    message: CustomMessage,
     buffers?: readonly DataView[],
   ) => void;
 }
@@ -135,8 +51,6 @@ export function getMarimoInternal<T extends ModelState>(
   return model[marimoSymbol];
 }
 
-export const MODEL_MANAGER = new ModelManager();
-
 export class Model<T extends ModelState> implements AnyModel<T> {
   #ANY_CHANGE_EVENT = "change";
   #dirtyFields: Map<keyof T, unknown>;
@@ -144,7 +58,14 @@ export class Model<T extends ModelState> implements AnyModel<T> {
   #comm: MarimoComm<T>;
   #listeners: Record<string, Set<EventHandler> | undefined> = {};
 
-  static _modelManager: ModelManager = MODEL_MANAGER;
+  /**
+   * Resolves models for the legacy ipywidgets escape hatch
+   * (`widget_manager.get_model`). Assigned by the registry module at
+   * import time to avoid an import cycle.
+   */
+  static _modelManager: {
+    getModel(model_id: WidgetModelId): Promise<Model<any>>;
+  };
 
   constructor(data: T, comm: MarimoComm<T>, signal?: AbortSignal) {
     this.#data = data;
@@ -164,9 +85,8 @@ export class Model<T extends ModelState> implements AnyModel<T> {
    */
   [marimoSymbol]: MarimoInternalApi<T> = {
     updateAndEmitDiffs: (value: T) => this.#updateAndEmitDiffs(value),
-    reemitState: () => this.#reemitState(),
     emitCustomMessage: (
-      message: Extract<AnyWidgetMessage, { method: "custom" }>,
+      message: CustomMessage,
       buffers?: readonly DataView[],
     ) => this.#emitCustomMessage(message, buffers),
   };
@@ -204,7 +124,7 @@ export class Model<T extends ModelState> implements AnyModel<T> {
     async get_model<TT extends ModelState>(
       model_id: WidgetModelId,
     ): Promise<AnyModel<TT>> {
-      const model = await Model._modelManager.get(model_id);
+      const model = await Model._modelManager.getModel(model_id);
       if (!model) {
         throw new Error(
           `Model not found with id: ${model_id}. This is likely because the model was not registered.`,
@@ -239,19 +159,43 @@ export class Model<T extends ModelState> implements AnyModel<T> {
     this.#comm.sendUpdate(partialData);
   }
 
-  on(eventName: string, callback: EventHandler): void {
+  /**
+   * Register an event listener.
+   *
+   * Passing `signal` ties the listener's lifetime to an `AbortSignal` —
+   * when the signal aborts, the listener is auto-removed. This is the
+   * preferred cleanup mechanism (matches the web platform pattern of
+   * `addEventListener({ signal })`) and is what `modelProxy` uses to
+   * scope listeners to view / binding lifetimes.
+   */
+  on(
+    eventName: string,
+    callback: EventHandler,
+    options?: { signal?: AbortSignal },
+  ): void {
+    const signal = options?.signal;
+    if (signal?.aborted) {
+      return;
+    }
     if (!this.#listeners[eventName]) {
       this.#listeners[eventName] = new Set();
     }
     this.#listeners[eventName].add(callback);
+    signal?.addEventListener("abort", () => this.off(eventName, callback), {
+      once: true,
+    });
   }
 
   #emit<K extends keyof T>(event: `change:${K & string}`, value: T[K]) {
-    if (!this.#listeners[event]) {
+    const listeners = this.#listeners[event];
+    if (!listeners) {
       return;
     }
-    const listeners = this.#listeners[event];
-    for (const listener of listeners) {
+    // Snapshot before iterating: a callback may unregister itself
+    // (typical signal-based abort handler calling `model.off`), and a
+    // Set mutated mid-iteration drops the very next element.
+    // oxlint-disable-next-line no-useless-spread -- snapshot is intentional
+    for (const listener of [...listeners]) {
       try {
         listener(value);
       } catch (error) {
@@ -274,29 +218,21 @@ export class Model<T extends ModelState> implements AnyModel<T> {
     });
   }
 
-  #reemitState() {
-    for (const [key, value] of Object.entries(this.#data) as [
-      keyof T & string,
-      T[keyof T],
-    ][]) {
-      this.#emit(`change:${key}`, value);
-    }
-    this.#emitAnyChange();
-  }
-
   /**
    * When receiving a message from the backend.
    * We want to notify all listeners with `msg:custom`
    */
   #emitCustomMessage(
-    message: Extract<AnyWidgetMessage, { method: "custom" }>,
+    message: CustomMessage,
     buffers: readonly DataView[] = [],
   ) {
     const listeners = this.#listeners["msg:custom"];
     if (!listeners) {
       return;
     }
-    for (const listener of listeners) {
+    // Snapshot before iterating: see `#emit` for rationale.
+    // oxlint-disable-next-line no-useless-spread -- snapshot is intentional
+    for (const listener of [...listeners]) {
       try {
         listener(message.content, buffers);
       } catch (error) {
@@ -311,7 +247,9 @@ export class Model<T extends ModelState> implements AnyModel<T> {
     if (!listeners) {
       return;
     }
-    for (const listener of listeners) {
+    // Snapshot before iterating: see `#emit` for rationale.
+    // oxlint-disable-next-line no-useless-spread -- snapshot is intentional
+    for (const listener of [...listeners]) {
       try {
         listener();
       } catch (error) {
@@ -320,121 +258,3 @@ export class Model<T extends ModelState> implements AnyModel<T> {
     }
   }, 0);
 }
-
-/**
- * Handle an incoming model lifecycle notification from the backend.
- *
- * Messages are dispatched by method type:
- * - "open": Initialize a new model or update existing one with initial state
- * - "update": Update model state with new values
- * - "custom": Forward custom message to model listeners
- * - "close": Remove model from manager
- */
-export async function handleWidgetMessage(
-  modelManager: ModelManager,
-  notification: NotificationMessageData<"model-lifecycle">,
-): Promise<void> {
-  const modelId = notification.model_id as WidgetModelId;
-  const msg = notification.message;
-
-  // Decode base64 buffers to DataViews (present in open/update/custom messages)
-  const base64Buffers: Base64String[] = "buffers" in msg ? msg.buffers : [];
-  const buffers = base64Buffers.map(base64ToDataView);
-
-  switch (msg.method) {
-    case "open": {
-      const { state, buffer_paths = [] } = msg;
-      const stateWithBuffers = decodeFromWire({
-        state,
-        bufferPaths: buffer_paths,
-        buffers,
-      });
-
-      // Check if a model already exists (created by the plugin using model_id reference)
-      // If so, just update its state instead of creating a duplicate
-      const existingModel = modelManager.getSync(modelId);
-      if (existingModel) {
-        getMarimoInternal(existingModel).updateAndEmitDiffs(stateWithBuffers);
-        return;
-      }
-
-      modelManager.create(modelId, (signal) => {
-        // In static exports there is no kernel, so comm calls are no-ops.
-        const comm: MarimoComm<ModelState> = isStaticNotebook()
-          ? {
-              sendUpdate: async () => undefined,
-              sendCustomMessage: async () => undefined,
-            }
-          : {
-              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),
-                });
-              },
-            };
-
-        return new Model(stateWithBuffers, comm, signal);
-      });
-      return;
-    }
-
-    case "custom": {
-      const model = await modelManager.get(modelId);
-      // For custom messages, we need to reconstruct the AnyWidgetMessage format
-      getMarimoInternal(model).emitCustomMessage(
-        { method: "custom", content: msg.content },
-        buffers,
-      );
-      return;
-    }
-
-    case "close":
-      BINDING_MANAGER.destroy(modelId);
-      modelManager.delete(modelId); // aborts the model's signal, clearing listeners
-      return;
-
-    case "update": {
-      const { state, buffer_paths = [] } = msg;
-      const stateWithBuffers = decodeFromWire({
-        state,
-        bufferPaths: buffer_paths,
-        buffers,
-      });
-      const model = await modelManager.get(modelId);
-      getMarimoInternal(model).updateAndEmitDiffs(stateWithBuffers);
-      return;
-    }
-
-    default:
-      assertNever(msg);
-  }
-}
-
-repl(MODEL_MANAGER, "MODEL_MANAGER");
-
-export const visibleForTesting = {
-  ModelManager,
-};