From ee1e20b82297dd2e7fb2da3a38fed82d0a207660 Mon Sep 17 00:00:00 2001
From: Kiran Gadhave <kirangadhave2@gmail.com>
Date: Thu, 25 Jun 2026 11:26:05 -0700
Subject: [PATCH] feat(mpl): reuse interactive figure manager across cell
 reruns

Cache the FigureManagerWebAgg per figure (reference-counted, weakly held) so
re-running mo.mpl.interactive(fig) reuses the manager instead of rebuilding it,
preserving toolbar state and skipping the canvas teardown and mpl.js
re-handshake. The frontend keeps the rendered figure mounted and rebinds the
existing socket to the new comm.
---
 .../mpl-interactive/MplInteractivePlugin.tsx  | 253 +++++++++++-------
 .../__tests__/MplInteractivePlugin.test.tsx   | 155 ++++++++++-
 .../mpl-interactive/mpl-websocket-shim.ts     |  10 +
 .../_plugins/ui/_impl/from_mpl_interactive.py | 187 +++++++++++--
 .../ui/_impl/test_from_mpl_interactive.py     | 169 +++++++++++-
 5 files changed, 637 insertions(+), 137 deletions(-)

diff --git a/frontend/src/plugins/impl/mpl-interactive/MplInteractivePlugin.tsx b/frontend/src/plugins/impl/mpl-interactive/MplInteractivePlugin.tsx
index 6bfac55c7f7..7fd3d6461ec 100644
--- a/frontend/src/plugins/impl/mpl-interactive/MplInteractivePlugin.tsx
+++ b/frontend/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;
   },
diff --git a/frontend/src/plugins/impl/mpl-interactive/__tests__/MplInteractivePlugin.test.tsx b/frontend/src/plugins/impl/mpl-interactive/__tests__/MplInteractivePlugin.test.tsx
index 9f1321eb2f1..62bfcdb68de 100644
--- a/frontend/src/plugins/impl/mpl-interactive/__tests__/MplInteractivePlugin.test.tsx
+++ b/frontend/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();
+  });
+});
diff --git a/frontend/src/plugins/impl/mpl-interactive/mpl-websocket-shim.ts b/frontend/src/plugins/impl/mpl-interactive/mpl-websocket-shim.ts
index ed4ca0b49b8..5eb11ccdd9e 100644
--- a/frontend/src/plugins/impl/mpl-interactive/mpl-websocket-shim.ts
+++ b/frontend/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.
diff --git a/marimo/_plugins/ui/_impl/from_mpl_interactive.py b/marimo/_plugins/ui/_impl/from_mpl_interactive.py
index a12b3a7e027..744e9959232 100644
--- a/marimo/_plugins/ui/_impl/from_mpl_interactive.py
+++ b/marimo/_plugins/ui/_impl/from_mpl_interactive.py
@@ -4,6 +4,7 @@
 import base64
 import functools
 import io
+import weakref
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 from uuid import uuid4
@@ -25,6 +26,8 @@
 from marimo._types.ids import WidgetModelId
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from matplotlib.figure import Figure, SubFigure
 
 LOGGER = _loggers.marimo_logger()
@@ -54,6 +57,87 @@ def _disconnect_owner_callbacks(canvas: Any, owner: Any) -> None:
         registry.disconnect(cid)
 
 
+class _FigureManagerRegistry:
+    """Shared `FigureManagerWebAgg` per figure, reference-counted by live
+    `mpl_interactive` instances.
+
+    matplotlib assumes one canvas/manager per figure for the figure's
+    lifetime, so every `mpl_interactive` wrapping the same figure reuses one
+    manager (and thus one toolbar, zoom history, and shared callback
+    registry). The manager is destroyed only when its last consumer is gone.
+
+    Both maps are weak so neither pins a figure on its own:
+      * The manager is held weakly; each live `mpl_interactive` keeps it alive
+        via `self._figure_manager`, so the entry drops once all are disposed.
+      * The figure is the weak key of the refcount map, so a discarded figure
+        evicts its own count without manual cleanup.
+
+    Bookkeeping only: manager creation, dpi/size capture, and canvas rebinding
+    are matplotlib-specific and stay with the caller.
+    """
+
+    def __init__(self) -> None:
+        self._managers: weakref.WeakValueDictionary[int, Any] = (
+            weakref.WeakValueDictionary()
+        )
+        self._refcounts: weakref.WeakKeyDictionary[Any, int] = (
+            weakref.WeakKeyDictionary()
+        )
+        # The figure's pristine dpi/size, captured before any interactive
+        # resize, to restore when the shared manager is destroyed. Kept here
+        # rather than on the manager so it shares the manager's lifecycle and
+        # stays typed instead of stringly monkey-patched onto a foreign object.
+        self._original_geometry: weakref.WeakKeyDictionary[
+            Any, tuple[float, tuple[float, float]]
+        ] = weakref.WeakKeyDictionary()
+
+    def acquire(
+        self, figure: Any, factory: Callable[[], Any]
+    ) -> tuple[Any, bool]:
+        """Return `(manager, created)` and increment the figure's refcount.
+
+        Reuses the cached manager if present, otherwise builds one via
+        `factory` and caches it. `created` is True only on the fresh build.
+        """
+        key = id(figure)
+        manager = self._managers.get(key)
+        created = False
+        if manager is None:
+            manager = factory()
+            self._managers[key] = manager
+            created = True
+        self._refcounts[figure] = self._refcounts.get(figure, 0) + 1
+        return manager, created
+
+    def set_original_geometry(
+        self, figure: Any, dpi: float, size_inches: tuple[float, float]
+    ) -> None:
+        self._original_geometry[figure] = (dpi, size_inches)
+
+    def original_geometry(
+        self, figure: Any
+    ) -> tuple[float, tuple[float, float]] | None:
+        return self._original_geometry.get(figure)
+
+    def release(self, figure: Any) -> bool:
+        """Decrement the figure's refcount.
+
+        Returns True when it reaches zero (the caller destroys the manager),
+        dropping every entry for the figure in that case.
+        """
+        count = self._refcounts.get(figure, 0) - 1
+        if count > 0:
+            self._refcounts[figure] = count
+            return False
+        self._refcounts.pop(figure, None)
+        self._managers.pop(id(figure), None)
+        self._original_geometry.pop(figure, None)
+        return True
+
+
+_FIGURE_MANAGERS = _FigureManagerRegistry()
+
+
 # Must match the className on the container div in MplInteractivePlugin.tsx
 _MPL_SCOPE = ".mpl-interactive-figure"
 
@@ -194,10 +278,27 @@ def __init__(self, figure: Figure | SubFigure) -> None:
         self._original_dpi = root.get_dpi()
         self._original_size_inches = tuple(root.get_size_inches())
 
-        # Create FigureManagerWebAgg
-        self._figure_manager = new_figure_manager_given_figure(
-            id(figure), figure
+        # Reuse the figure's manager across reruns so toolbar state, zoom
+        # history, and the shared callback registry survive. Only the comm is
+        # recreated per element (below).
+        manager, created = _FIGURE_MANAGERS.acquire(
+            figure,
+            lambda: new_figure_manager_given_figure(id(figure), figure),
         )
+        if created:
+            # Capture the pristine dpi/size once, before any interactive
+            # resize mutates them; restored when the shared manager is
+            # finally destroyed (refcount reaches zero).
+            _FIGURE_MANAGERS.set_original_geometry(
+                figure, self._original_dpi, self._original_size_inches
+            )
+        elif root.canvas is not manager.canvas:
+            # matplotlib >= 3.11 resets `figure.canvas` to a bare
+            # FigureCanvasBase on close, and marimo closes figures after every
+            # cell run (`close_figures()` -> `plt.close("all")`). Re-bind the
+            # cached canvas so the first frame after a rerun renders against it.
+            root.set_canvas(manager.canvas)
+        self._figure_manager = manager
 
         # Get figure dimensions in CSS pixels for initial sizing.
         # get_width_height() returns device pixels (includes DPI scaling).
@@ -265,10 +366,9 @@ def _initialize(
             ctx.cell_lifecycle_registry.add(
                 _MplCleanupHandle(
                     comm=self._comm,
+                    figure=self._figure,
                     figure_manager=self._figure_manager,
                     sync_ws=self._sync_ws,
-                    original_dpi=self._original_dpi,
-                    original_size_inches=self._original_size_inches,
                 )
             )
         except ContextNotInitializedError:
@@ -354,17 +454,26 @@ def _mime_(self) -> tuple[KnownMimeType, str]:
 
 
 class _MplCleanupHandle(CellLifecycleItem):
-    """Cleans up the matplotlib figure manager and MarimoComm on cell re-run or deletion."""
+    """Cleans up an `mpl_interactive` on cell re-run or deletion.
+
+    The comm and websocket are per-element and always torn down. The figure
+    manager is shared across every element wrapping the same figure, so it is
+    destroyed only when the last of those elements is disposed (its refcount
+    reaches zero). Callers that do not pass a `figure` fall back to destroying
+    the manager unconditionally.
+    """
 
     def __init__(
         self,
         comm: MarimoComm,
-        original_dpi: float,
-        original_size_inches: tuple[float, float],
+        figure: Any = None,
         figure_manager: Any = None,
         sync_ws: Any = None,
+        original_dpi: float | None = None,
+        original_size_inches: tuple[float, float] | None = None,
     ) -> None:
         self._comm = comm
+        self._figure = figure
         self._figure_manager = figure_manager
         self._sync_ws = sync_ws
         self._original_dpi = original_dpi
@@ -376,32 +485,52 @@ def create(self, context: RuntimeContext) -> None:
     def dispose(self, context: RuntimeContext, deletion: bool) -> bool:
         del context, deletion
         fm = self._figure_manager
-        if fm is not None:
-            if self._sync_ws is not None:
-                try:
-                    fm.remove_web_socket(self._sync_ws)
-                except Exception:
-                    LOGGER.exception("Failed to remove mpl web socket")
-
-            canvas = fm.canvas
-            toolbar = getattr(canvas, "toolbar", None)
-            if toolbar is not None:
-                try:
-                    _disconnect_owner_callbacks(canvas, toolbar)
-                except Exception:
-                    LOGGER.exception(
-                        "Failed to disconnect mpl toolbar callbacks"
-                    )
+        if fm is not None and self._sync_ws is not None:
+            try:
+                fm.remove_web_socket(self._sync_ws)
+            except Exception:
+                LOGGER.exception("Failed to remove mpl web socket")
 
+        if self._figure is None:
+            # No shared refcount (manual/legacy callers): destroy immediately,
+            # restoring the geometry passed at construction.
+            self._destroy_manager(
+                self._original_dpi, self._original_size_inches
+            )
+        else:
+            # Read the captured geometry before release drops it.
+            geometry = _FIGURE_MANAGERS.original_geometry(self._figure)
+            if _FIGURE_MANAGERS.release(self._figure):
+                dpi, size = geometry if geometry is not None else (None, None)
+                self._destroy_manager(dpi, size)
+
+        self._comm.close()
+        return True
+
+    def _destroy_manager(
+        self,
+        dpi: float | None,
+        size_inches: tuple[float, float] | None,
+    ) -> None:
+        fm = self._figure_manager
+        if fm is None:
+            return
+
+        canvas = fm.canvas
+        toolbar = getattr(canvas, "toolbar", None)
+        if toolbar is not None:
+            try:
+                _disconnect_owner_callbacks(canvas, toolbar)
+            except Exception:
+                LOGGER.exception("Failed to disconnect mpl toolbar callbacks")
+
+        if dpi is not None and size_inches is not None:
             try:
                 # get the root figure (in case of Subfigure) which handles dpi
                 root = canvas.figure.figure
-                root.set_dpi(self._original_dpi)
-                root.set_size_inches(*self._original_size_inches)
+                root.set_dpi(dpi)
+                root.set_size_inches(*size_inches)
             except Exception:
                 LOGGER.exception(
                     "Failed to restore mpl figure dpi/size on dispose"
                 )
-
-        self._comm.close()
-        return True
diff --git a/tests/_plugins/ui/_impl/test_from_mpl_interactive.py b/tests/_plugins/ui/_impl/test_from_mpl_interactive.py
index a63fb4afe40..488d57696e7 100644
--- a/tests/_plugins/ui/_impl/test_from_mpl_interactive.py
+++ b/tests/_plugins/ui/_impl/test_from_mpl_interactive.py
@@ -225,13 +225,14 @@ def test_dpi_does_not_compound_across_reruns(self) -> None:
 
             # Simulate cell teardown — the cleanup handle is what marimo
             # registers via cell_lifecycle_registry; running it directly
-            # avoids needing a live runtime context in the test.
+            # avoids needing a live runtime context in the test. Passing the
+            # figure drives the refcount to zero so the shared manager is
+            # destroyed and the next iteration builds a fresh one.
             cleanup = _MplCleanupHandle(
                 comm=element._comm,
+                figure=fig,
                 figure_manager=element._figure_manager,
                 sync_ws=element._sync_ws,
-                original_dpi=element._original_dpi,
-                original_size_inches=element._original_size_inches,
             )
             cleanup.dispose(context=MagicMock(), deletion=False)
 
@@ -467,19 +468,22 @@ def test_dispose_disconnects_toolbar_callbacks(self) -> None:
         assert self._count_toolbar_handlers(fig, "button_release_event") == 1
 
         # Simulate cell teardown: marimo's runtime invokes dispose on the
-        # lifecycle handle when the cell re-runs or is deleted.
+        # lifecycle handle when the cell re-runs or is deleted. Passing the
+        # figure drives the refcount to zero, which destroys the shared
+        # manager and disconnects its toolbar from the registry.
         cleanup = _MplCleanupHandle(
             comm=elem1._comm,
+            figure=fig,
             figure_manager=elem1._figure_manager,
             sync_ws=elem1._sync_ws,
-            original_dpi=elem1._original_dpi,
-            original_size_inches=elem1._original_size_inches,
         )
         cleanup.dispose(context=MagicMock(), deletion=False)
+        assert self._count_toolbar_handlers(fig, "button_press_event") == 0
+        assert self._count_toolbar_handlers(fig, "button_release_event") == 0
 
-        # Simulate the cell re-running: a new mpl_interactive on the same
-        # figure object. After dispose, only the new toolbar's callbacks
-        # should be live on the shared registry.
+        # Re-running the cell builds a fresh manager (the cache entry was
+        # dropped at zero), so exactly one toolbar's callbacks are live —
+        # never stacked.
         with patch("marimo._plugins.ui._impl.comm.broadcast_notification"):
             _elem2 = mpl_interactive(fig)
 
@@ -489,6 +493,153 @@ def test_dispose_disconnects_toolbar_callbacks(self) -> None:
         plt.close(fig)
 
 
+@pytest.mark.requires("matplotlib")
+class TestFigureManagerReuse:
+    """A figure's manager is shared across reruns and cells, and destroyed
+    only when its last consumer is disposed.
+
+    matplotlib assumes one canvas/manager per figure for the figure's
+    lifetime, so wrapping the same figure again must reuse the existing
+    manager (preserving toolbar state) rather than building a new one.
+    """
+
+    @staticmethod
+    def _new_element(fig: Any) -> Any:
+        from marimo._plugins.ui._impl.from_mpl_interactive import (
+            mpl_interactive,
+        )
+
+        with patch("marimo._plugins.ui._impl.comm.broadcast_notification"):
+            return mpl_interactive(fig)
+
+    @staticmethod
+    def _cleanup_handle(element: Any, fig: Any) -> Any:
+        from marimo._plugins.ui._impl.from_mpl_interactive import (
+            _MplCleanupHandle,
+        )
+
+        return _MplCleanupHandle(
+            comm=element._comm,
+            figure=fig,
+            figure_manager=element._figure_manager,
+            sync_ws=element._sync_ws,
+        )
+
+    @staticmethod
+    def _toolbar_handler_count(fig: Any, signal: str) -> int:
+        handlers = fig.canvas.callbacks.callbacks.get(signal, {})
+        count = 0
+        for ref in handlers.values():
+            fn = ref() if callable(ref) else ref
+            if type(getattr(fn, "__self__", None)).__name__ == (
+                "NavigationToolbar2WebAgg"
+            ):
+                count += 1
+        return count
+
+    def test_reuse_shares_manager_and_preserves_toolbar(self) -> None:
+        import matplotlib.pyplot as plt
+
+        fig, ax = plt.subplots()
+        ax.plot([1, 2, 3])
+
+        elem1 = self._new_element(fig)
+        manager = elem1._figure_manager
+
+        # A toolbar mode set on the first wrapper must survive the second.
+        elem1._figure_manager.canvas.toolbar.pan()
+        mode = str(elem1._figure_manager.canvas.toolbar.mode)
+
+        elem2 = self._new_element(fig)
+
+        assert elem2._figure_manager is manager
+        assert str(elem2._figure_manager.canvas.toolbar.mode) == mode
+
+        plt.close(fig)
+
+    def test_different_figure_gets_new_manager(self) -> None:
+        import matplotlib.pyplot as plt
+
+        fig1, ax1 = plt.subplots()
+        ax1.plot([1, 2, 3])
+        fig2, ax2 = plt.subplots()
+        ax2.plot([3, 2, 1])
+
+        elem1 = self._new_element(fig1)
+        elem2 = self._new_element(fig2)
+
+        assert elem1._figure_manager is not elem2._figure_manager
+
+        plt.close(fig1)
+        plt.close(fig2)
+
+    def test_manager_destroyed_when_last_consumer_disposes(self) -> None:
+        import matplotlib.pyplot as plt
+
+        from marimo._plugins.ui._impl.from_mpl_interactive import (
+            _FIGURE_MANAGERS,
+        )
+
+        fig, ax = plt.subplots()
+        ax.plot([1, 2, 3])
+
+        elem = self._new_element(fig)
+        assert _FIGURE_MANAGERS._managers.get(id(fig)) is elem._figure_manager
+        assert self._toolbar_handler_count(fig, "button_press_event") == 1
+
+        self._cleanup_handle(elem, fig).dispose(
+            context=MagicMock(), deletion=False
+        )
+
+        assert _FIGURE_MANAGERS._managers.get(id(fig)) is None
+        assert fig not in _FIGURE_MANAGERS._refcounts
+        assert self._toolbar_handler_count(fig, "button_press_event") == 0
+
+        plt.close(fig)
+
+    def test_two_cells_share_manager_independent_dispose(self) -> None:
+        import matplotlib.pyplot as plt
+
+        from marimo._plugins.ui._impl.from_mpl_interactive import (
+            _FIGURE_MANAGERS,
+        )
+
+        fig, ax = plt.subplots()
+        ax.plot([1, 2, 3])
+
+        elem1 = self._new_element(fig)
+        elem2 = self._new_element(fig)
+        manager = elem1._figure_manager
+
+        assert elem2._figure_manager is manager
+        assert _FIGURE_MANAGERS._refcounts[fig] == 2
+        assert elem1._sync_ws in manager.web_sockets
+        assert elem2._sync_ws in manager.web_sockets
+
+        # Disposing one consumer leaves the manager alive for the other.
+        self._cleanup_handle(elem1, fig).dispose(
+            context=MagicMock(), deletion=False
+        )
+
+        assert _FIGURE_MANAGERS._managers.get(id(fig)) is manager
+        assert _FIGURE_MANAGERS._refcounts[fig] == 1
+        assert elem1._sync_ws not in manager.web_sockets
+        assert elem2._sync_ws in manager.web_sockets
+        # The shared toolbar's callbacks stay connected until the last consumer.
+        assert self._toolbar_handler_count(fig, "button_press_event") == 1
+
+        # Disposing the second fully tears down.
+        self._cleanup_handle(elem2, fig).dispose(
+            context=MagicMock(), deletion=False
+        )
+
+        assert _FIGURE_MANAGERS._managers.get(id(fig)) is None
+        assert fig not in _FIGURE_MANAGERS._refcounts
+        assert self._toolbar_handler_count(fig, "button_press_event") == 0
+
+        plt.close(fig)
+
+
 @pytest.mark.requires("matplotlib")
 class TestMplInteractiveArgs:
     """Test that mpl_interactive passes correct args to the UIElement."""