@marimo-team/islands 0.23.12-dev5 → 0.23.12-dev7

package/src/plugins/impl/mpl-interactive/MplInteractivePlugin.tsx

@@ -174,99 +174,99 @@ const MplInteractiveSlot = (props: IPluginProps<ModelIdRef, Data>) => {
   const containerRef = useRef<HTMLDivElement>(null);
   const figureRef = useRef<MplFigure | null>(null);
   const wsRef = useRef<MplCommWebSocket | null>(null);
+  // Sends to the currently bound backend model. Re-pointed on every (re)bind
+  // so the persistent socket and toolbar downloads always reach the live comm.
+  const sendRef = useRef<(msg: unknown) => void>(Functions.NOOP);
+  // Detaches the model bound by the most recent bindModel call. Shared between
+  // the mount and rebind effects so a rerun disposes the prior model's
+  // listener before attaching the next one (never stacking listeners).
+  const boundModelCleanupRef = useRef<(() => void) | undefined>(undefined);
+  // Latest model id, read by the mount effect without being a dependency:
+  // the figure is built once, and switching to a new model is the rebind
+  // effect's job, not a reason to tear the canvas down.
+  const modelIdRef = useRef(modelId);
+  modelIdRef.current = modelId;
+  // The data attributes are re-parsed into fresh objects on every rerun, so
+  // `toolbarImages` changes identity even when its contents do not. Read it
+  // from a ref so it can't retrigger the mount effect and rebuild the canvas.
+  const toolbarImagesRef = useRef(toolbarImages);
+  toolbarImagesRef.current = toolbarImages;
+
+  // Bind the already-rendered figure/socket to a backend model, leaving the
+  // DOM, figure, and socket in place so they survive across reruns. Disposes
+  // the previously bound model first and records the new cleanup in
+  // boundModelCleanupRef, so exactly one model listener is ever attached.
+  const bindModel = useCallback(async (id: WidgetModelId): Promise<void> => {
+    const fakeWs = wsRef.current;
+    if (!fakeWs) {
+      return;
+    }
 
-  const setupFigure = useCallback(
-    async (container: HTMLElement) => {
-      // Load mpl.js globally (only once, via <script src>)
-      await ensureMplJs(mplJsUrl);
+    let model: Model<ModelState>;
+    try {
+      model = await MODEL_MANAGER.get(id);
+    } catch {
+      Logger.error("Failed to get model for mpl interactive", id);
+      return;
+    }
 
-      if (!window.mpl) {
-        Logger.error("mpl.js failed to load");
-        return;
-      }
+    // The figure may have been torn down (unmount, or a structural rebuild)
+    // while we awaited the model; don't wire a listener that would outlive it.
+    if (wsRef.current !== fakeWs) {
+      return;
+    }
 
-      // Get the model from MODEL_MANAGER
-      let model: Model<ModelState>;
-      try {
-        model = await MODEL_MANAGER.get(modelId);
-      } catch {
-        Logger.error("Failed to get model for mpl interactive", modelId);
+    // Detach the previously bound model before wiring the new one.
+    boundModelCleanupRef.current?.();
+
+    // Re-point outbound traffic at this model without recreating the socket,
+    // so mpl.js's onopen/onmessage wiring stays intact.
+    const send = (msg: unknown) => model.send(msg);
+    fakeWs.setSendHandler(send);
+    sendRef.current = send;
+
+    // Listen for backend → frontend messages via model custom events
+    const handleCustomMessage = (
+      content: { type: string; data?: unknown; format?: string },
+      buffers?: readonly DataView[],
+    ) => {
+      if (!content) {
         return;
       }
 
-      // Create the fake WebSocket
-      const fakeWs = new MplCommWebSocket((msg: unknown) => {
-        // Send from frontend → backend via model custom message
-        model.send(msg);
-      });
-      wsRef.current = fakeWs;
-
-      // Listen for backend → frontend messages via model custom events
-      const handleCustomMessage = (
-        content: { type: string; data?: unknown; format?: string },
-        buffers?: readonly DataView[],
-      ) => {
-        if (!content) {
-          return;
-        }
-
-        if (content.type === "json") {
-          fakeWs.receiveJson(content.data);
-        } else if (content.type === "binary" && buffers && buffers.length > 0) {
-          fakeWs.receiveBinary(buffers[0]);
-        } else if (
-          content.type === "download" &&
-          buffers &&
-          buffers.length > 0
-        ) {
-          const fmt = content.format || "png";
-          const dv = buffers[0];
-          const ab = dv.buffer.slice(
-            dv.byteOffset,
-            dv.byteOffset + dv.byteLength,
-          ) as ArrayBuffer;
-          downloadBlob(
-            new Blob([ab], { type: `image/${fmt}` }),
-            `figure.${fmt}`,
-          );
-        }
-      };
+      if (content.type === "json") {
+        fakeWs.receiveJson(content.data);
+      } else if (content.type === "binary" && buffers && buffers.length > 0) {
+        fakeWs.receiveBinary(buffers[0]);
+      } else if (content.type === "download" && buffers && buffers.length > 0) {
+        const fmt = content.format || "png";
+        const dv = buffers[0];
+        const ab = dv.buffer.slice(
+          dv.byteOffset,
+          dv.byteOffset + dv.byteLength,
+        ) as ArrayBuffer;
+        downloadBlob(new Blob([ab], { type: `image/${fmt}` }), `figure.${fmt}`);
+      }
+    };
 
-      model.on("msg:custom", handleCustomMessage as any);
+    model.on("msg:custom", handleCustomMessage as any);
 
-      // Create the mpl figure
-      const figId = modelId;
-      const ondownload = (_figure: MplFigure, format: string) => {
-        // Send download request to backend
-        model.send({ type: "download", format });
-      };
+    // Replay the mpl.js handshake against the new comm so the backend
+    // re-establishes image mode and pushes a frame. The figure DOM and the
+    // backend manager's toolbar state are left untouched.
+    fakeWs.onopen?.();
 
-      const fig = new window.mpl.figure(figId, fakeWs, ondownload, container);
-      figureRef.current = fig;
-
-      // Set the canvas_div to the backend's figure size so the
-      // ResizeObserver doesn't trigger an immediate resize cycle.
-      // mpl.js creates: fig.root > [titlebar, canvas_div, toolbar]
-      const canvasDiv = fig.root.querySelector<HTMLElement>("div[tabindex]");
-      if (canvasDiv) {
-        canvasDiv.style.width = `${width}px`;
-        canvasDiv.style.height = `${height}px`;
+    boundModelCleanupRef.current = () => {
+      model.off("msg:custom", handleCustomMessage as any);
+      if (sendRef.current === send) {
+        sendRef.current = Functions.NOOP;
       }
+    };
+  }, []);
 
-      // Trigger the onopen callback to start communication
-      // mpl.js sends initial messages in onopen
-      setTimeout(() => {
-        fakeWs.onopen?.();
-      }, 0);
-
-      return () => {
-        model.off("msg:custom", handleCustomMessage as any);
-        fakeWs.close();
-      };
-    },
-    [modelId, mplJsUrl, width, height],
-  );
-
+  // Mount: build the DOM, mpl figure, and socket once. modelId is read from a
+  // ref and intentionally omitted from the deps — rebinding to a new model is
+  // handled by the effect below, not by rebuilding the canvas.
   useEffect(() => {
     const container = containerRef.current;
     if (!container) {
@@ -280,34 +280,90 @@ const MplInteractiveSlot = (props: IPluginProps<ModelIdRef, Data>) => {
     const removeCss = injectCss(container, cssUrl);
 
     // Patch toolbar images
-    const removeImageObserver = patchToolbarImages(container, toolbarImages);
+    const removeImageObserver = patchToolbarImages(
+      container,
+      toolbarImagesRef.current,
+    );
 
-    let cleanup: (() => void) | undefined;
     let cancelled = false;
 
-    setupFigure(container)
-      .then((cleanupFn) => {
-        if (cancelled) {
-          cleanupFn?.();
-          return;
-        }
-        cleanup = cleanupFn;
-      })
-      .catch((error) => {
-        if (!cancelled) {
-          Logger.error("Failed to set up MPL interactive figure", error);
-        }
+    const setup = async () => {
+      // Load mpl.js globally (only once, via <script src>)
+      await ensureMplJs(mplJsUrl);
+
+      if (!window.mpl) {
+        Logger.error("mpl.js failed to load");
+        return;
+      }
+
+      // The send handler is swapped per bind; route through sendRef so the
+      // socket outlives any single model.
+      const fakeWs = new MplCommWebSocket((msg: unknown) => {
+        sendRef.current(msg);
       });
+      wsRef.current = fakeWs;
+
+      const ondownload = (_figure: MplFigure, format: string) => {
+        sendRef.current({ type: "download", format });
+      };
+
+      const fig = new window.mpl.figure(
+        modelIdRef.current,
+        fakeWs,
+        ondownload,
+        container,
+      );
+      figureRef.current = fig;
+
+      // Set the canvas_div to the backend's figure size so the
+      // ResizeObserver doesn't trigger an immediate resize cycle.
+      // mpl.js creates: fig.root > [titlebar, canvas_div, toolbar]
+      const canvasDiv = fig.root.querySelector<HTMLElement>("div[tabindex]");
+      if (canvasDiv) {
+        canvasDiv.style.width = `${width}px`;
+        canvasDiv.style.height = `${height}px`;
+      }
+
+      await bindModel(modelIdRef.current);
+    };
+
+    setup().catch((error) => {
+      if (!cancelled) {
+        Logger.error("Failed to set up MPL interactive figure", error);
+      }
+    });
 
     return () => {
       cancelled = true;
+      boundModelCleanupRef.current?.();
+      boundModelCleanupRef.current = undefined;
       removeCss();
       removeImageObserver();
-      cleanup?.();
+      wsRef.current?.close();
+      wsRef.current = null;
+      figureRef.current = null;
       // Clear DOM on unmount so stale content doesn't linger
       container.innerHTML = "";
     };
-  }, [modelId, cssUrl, toolbarImages, setupFigure]);
+  }, [mplJsUrl, cssUrl, width, height, bindModel]);
+
+  // Rebind to a new model when the cell re-runs, keeping the rendered figure
+  // and toolbar in place. The initial bind is owned by the mount effect, so
+  // skip the first run here.
+  const isInitialBindRef = useRef(true);
+  useEffect(() => {
+    if (isInitialBindRef.current) {
+      isInitialBindRef.current = false;
+      return;
+    }
+
+    // bindModel disposes the previously bound model and guards against a
+    // teardown that races the awaited model lookup, so no per-run cleanup is
+    // needed here; the mount effect's cleanup detaches the final bind.
+    bindModel(modelId).catch((error) => {
+      Logger.error("Failed to rebind MPL interactive figure", error);
+    });
+  }, [modelId, bindModel]);
 
   // Re-request figure when tab becomes visible
   useEventListener(document, "visibilitychange", () => {
@@ -324,6 +380,7 @@ const MplInteractiveSlot = (props: IPluginProps<ModelIdRef, Data>) => {
 export const visibleForTesting = {
   ensureMplJs,
   injectCss,
+  MplInteractiveSlot,
   resetMplJsLoading: () => {
     mplJsLoading = null;
   },

package/src/plugins/impl/mpl-interactive/__tests__/MplInteractivePlugin.test.tsx

@@ -1,4 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
+/* oxlint-disable typescript/no-explicit-any */
+/* oxlint-disable marimo/prefer-object-params -- the mocked mpl.js figure
+   constructor must match its real positional signature. */
+import { render, waitFor } from "@testing-library/react";
 import type { ExtractAtomValue } from "jotai";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { hasRunAnyCellAtom } from "@/components/editor/cell/useRunCells";
@@ -7,9 +11,12 @@ import { parseUserConfig } from "@/core/config/config-schema";
 import { initialModeAtom } from "@/core/mode";
 import { store } from "@/core/state/jotai";
 import { Logger } from "@/utils/Logger";
+import { MODEL_MANAGER, Model } from "@/plugins/impl/anywidget/model";
+import type { WidgetModelId } from "@/plugins/impl/anywidget/types";
 import { visibleForTesting } from "../MplInteractivePlugin";
 
-const { ensureMplJs, injectCss, resetMplJsLoading } = visibleForTesting;
+const { ensureMplJs, injectCss, MplInteractiveSlot, resetMplJsLoading } =
+  visibleForTesting;
 
 /**
  * Clear every "notebook trust" signal `isTrustedVirtualFileUrl` consults so
@@ -151,3 +158,149 @@ describe("MplInteractivePlugin URL validation", () => {
     });
   });
 });
+
+const asModelId = (id: string): WidgetModelId => id as WidgetModelId;
+
+/**
+ * Minimal stand-in for the global `window.mpl.figure` constructor. mpl.js
+ * builds the canvas DOM and wires `onopen`/`onmessage` onto the socket at
+ * construction; the mock reproduces just enough of that for the slot to mount
+ * and rebind.
+ */
+function installMplFigureMock(): ReturnType<typeof vi.fn> {
+  const ctor = vi.fn(function (
+    this: any,
+    id: string,
+    ws: any,
+    _ondownload: unknown,
+    container: HTMLElement,
+  ) {
+    this.id = id;
+    this.ws = ws;
+    const root = document.createElement("div");
+    const canvasDiv = document.createElement("div");
+    canvasDiv.setAttribute("tabindex", "0");
+    root.append(canvasDiv);
+    container.append(root);
+    this.root = root;
+    this.send_message = vi.fn();
+    ws.onopen = vi.fn();
+    ws.onmessage = vi.fn();
+  });
+  (window as unknown as { mpl: unknown }).mpl = {
+    figure: ctor,
+    toolbar_items: [],
+  };
+  return ctor;
+}
+
+function makeModel(): Model<Record<string, never>> {
+  return new Model(
+    {},
+    {
+      sendUpdate: vi.fn().mockResolvedValue(undefined),
+      sendCustomMessage: vi.fn().mockResolvedValue(undefined),
+    },
+  );
+}
+
+function makeProps(modelId: WidgetModelId) {
+  return {
+    data: {
+      mplJsUrl: "./@file/1-mpl.js",
+      cssUrl: "./@file/2-mpl.css",
+      toolbarImages: {},
+      width: 640,
+      height: 480,
+    },
+    value: { model_id: modelId },
+    host: document.createElement("div"),
+    setValue: vi.fn(),
+    functions: {},
+  } as any;
+}
+
+describe("MplInteractiveSlot rerun rebinding", () => {
+  beforeEach(() => {
+    vi.spyOn(Logger, "error").mockImplementation(() => {});
+    resetMplJsLoading();
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+    delete (window as { mpl?: unknown }).mpl;
+  });
+
+  it("rebinds to a new model without rebuilding the figure DOM", async () => {
+    const ctor = installMplFigureMock();
+    const idA = asModelId("model-a");
+    const idB = asModelId("model-b");
+    MODEL_MANAGER.set(idA, makeModel());
+    MODEL_MANAGER.set(idB, makeModel());
+
+    const { container, rerender } = render(
+      <MplInteractiveSlot {...makeProps(idA)} />,
+    );
+
+    await waitFor(() => expect(ctor).toHaveBeenCalledTimes(1));
+
+    const figureRoot = ctor.mock.instances[0].root as HTMLElement;
+    const socket = ctor.mock.calls[0][1];
+    const slot = container.querySelector(".mpl-interactive-figure");
+    expect(slot?.contains(figureRoot)).toBe(true);
+    // One handshake on the initial bind.
+    expect(socket.onopen).toHaveBeenCalledTimes(1);
+    const setSendHandler = vi.spyOn(socket, "setSendHandler");
+
+    // Cell re-run: only the model id changes.
+    rerender(<MplInteractiveSlot {...makeProps(idB)} />);
+
+    // The new model is bound through the existing socket, with a fresh
+    // handshake, and the figure is never reconstructed.
+    await waitFor(() => expect(socket.onopen).toHaveBeenCalledTimes(2));
+    expect(ctor).toHaveBeenCalledTimes(1);
+    expect(setSendHandler).toHaveBeenCalledTimes(1);
+    // The same rendered DOM is still in place — not cleared and rebuilt.
+    expect(container.querySelector(".mpl-interactive-figure")).toBe(slot);
+    expect(slot?.contains(figureRoot)).toBe(true);
+  });
+
+  it("detaches the previous model's listener on each rerun (no buildup)", async () => {
+    installMplFigureMock();
+    // Unique ids: MODEL_MANAGER is a module singleton whose deferreds resolve
+    // once, so reusing ids from another test would return that test's models.
+    const idA = asModelId("leak-a");
+    const idB = asModelId("leak-b");
+    const idC = asModelId("leak-c");
+    const modelA = makeModel();
+    const modelB = makeModel();
+    const modelC = makeModel();
+    MODEL_MANAGER.set(idA, modelA);
+    MODEL_MANAGER.set(idB, modelB);
+    MODEL_MANAGER.set(idC, modelC);
+
+    const onA = vi.spyOn(modelA, "on");
+    const offA = vi.spyOn(modelA, "off");
+    const onB = vi.spyOn(modelB, "on");
+    const offB = vi.spyOn(modelB, "off");
+    const onC = vi.spyOn(modelC, "on");
+    const offC = vi.spyOn(modelC, "off");
+
+    const { rerender } = render(<MplInteractiveSlot {...makeProps(idA)} />);
+    await waitFor(() =>
+      expect(onA).toHaveBeenCalledWith("msg:custom", expect.any(Function)),
+    );
+
+    rerender(<MplInteractiveSlot {...makeProps(idB)} />);
+    await waitFor(() => expect(onB).toHaveBeenCalled());
+
+    rerender(<MplInteractiveSlot {...makeProps(idC)} />);
+    await waitFor(() => expect(onC).toHaveBeenCalled());
+
+    // Each superseded model had its listener detached exactly once when the
+    // next bind replaced it; the current model's listener stays attached.
+    expect(offA).toHaveBeenCalledTimes(1);
+    expect(offB).toHaveBeenCalledTimes(1);
+    expect(offC).not.toHaveBeenCalled();
+  });
+});

package/src/plugins/impl/mpl-interactive/mpl-websocket-shim.ts

@@ -21,6 +21,16 @@ export class MplCommWebSocket {
     this.sendFn = sendFn;
   }
 
+  /**
+   * Re-point outbound sends at a new backend comm without recreating the
+   * socket. The figure manager is reused across cell reruns, so the same
+   * socket instance stays bound to mpl.js (which wires `onopen`/`onmessage`
+   * onto it at construction); only the comm behind it changes.
+   */
+  setSendHandler(sendFn: (msg: unknown) => void): void {
+    this.sendFn = sendFn;
+  }
+
   /**
    * Called by mpl.js to send a message to the backend.
    * mpl.js always sends JSON strings.