feat(mpl): reuse interactive figure manager across cell reruns

[kirangadhave]

📝 Summary

mo.mpl.interactive(fig) built a fresh FigureManagerWebAgg + canvas + toolbar on every cell rerun, even for the same figure. matplotlib assumes one canvas per figure for the figure's lifetime (the event-callback registry lives on the figure and is shared across its canvases), so rebuilding tore down toolbar state, re-handshook mpl.js with a visible flicker, and stacked handlers on the shared registry.

matplotlib's WebAgg backend is built around a single long-lived FigureManagerWebAgg per figure, with clients multiplexed onto it via add_web_socket/remove_web_socket (the manager keeps a web_sockets set, and add_web_socket resizes and refreshes the new client). The official embedding_webagg example marimo adapted from follows exactly this pattern. Rebuilding the manager per rerun and attaching a single socket to the throwaway inverts that design.

Cache the manager per figure in a reference-counted, weakly-held registry keyed by id(figure). Re-running mo.mpl.interactive(fig) reuses the manager and only recreates the per-element comm; multiple cells wrapping the same figure share one manager. The manager is destroyed (toolbar callbacks disconnected, dpi/size restored) only when the last consumer is disposed. The figure's pristine dpi/size lives in the registry rather than monkey-patched onto the manager.

On the frontend, a model-id change for the same figure no longer clears and rebuilds the canvas DOM. The mount effect builds the figure once; a separate rebind effect re-points the existing MplCommWebSocket at the new comm, so mpl.js's onopen/onmessage wiring stays intact and toolbar state survives.

Result: re-running a cell on the same figure preserves toolbar mode and zoom/pan history with no flicker or re-sync teardown, and two cells displaying one figure stay in sync through a single manager.

Closes MO-6172

Related (already patched separately; this PR removes the shared rebuild-every-rerun root cause, the defensive code is left in place):

• MO-6160 (fix(mpl): disconnect toolbar callbacks on cell rerun #9531): leaked _zoom_pan_handler on the figure-shared callback registry.
• MO-6171 (fix(mpl): contain comm callback errors in mpl_interactive #9532): comm-callback exceptions crashing the kernel subprocess.
• MO-6533 (fix: support matplotlib 3.11 interactive #9990): detached WebAgg canvas after close_figures() under matplotlib >= 3.11.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
marimo-docs	Ready	Preview, Comment	Jun 26, 2026 4:43am

[cubic-dev-ai]

No issues found across 5 files

Architecture diagram

sequenceDiagram
    participant Client as Frontend (React)
    participant MplSlot as MplInteractiveSlot
    participant MplComm as MplCommWebSocket
    participant Model as WidgetModel
    participant Backend as Python Kernel
    participant Registry as _FigureManagerRegistry
    participant Mgr as FigureManagerWebAgg
    participant Fig as matplotlib Figure

    Note over Client,Fig: INITIAL CELL EXECUTION
    
    Client->>MplSlot: mount (useEffect)
    MplSlot->>MplComm: new MplCommWebSocket(sendFn)
    MplSlot->>MplComm: create mpl.figure(id, ws, ondownload, container)
    MplSlot->>MplComm: trigger fakeWs.onopen()
    MplComm->>Model: model.send(msg) - handshake
    Model-->>MplComm: msg:custom events
    
    Note over Backend,Fig: Python side - manager acquisition
    
    Backend->>Registry: acquire(fig, factory)
    Registry->>Registry: check cache by id(fig)
    alt Cache miss (first time)
        Registry->>Mgr: factory() - new_figure_manager_given_figure()
        Mgr-->>Registry: manager
        Registry->>Registry: set_original_geometry(fig, dpi, size)
    else Cache hit (rerun)
        Registry-->>Backend: cached manager
        Backend->>Fig: root.set_canvas(manager.canvas) if needed
    end
    Registry-->>Backend: (manager, created)
    Backend->>Mgr: store as self._figure_manager
    
    Note over Backend: Create per-element comm & websocket
    
    Backend->>Mgr: add_web_socket(sync_ws)
    Mgr-->>Backend: websocket registered
    Backend->>Backend: register _MplCleanupHandle with ref to figure

    Note over Client,Fig: CELL RE-RUN (same figure)
    
    Client->>MplSlot: rerender with new modelId
    MplSlot->>MplSlot: read modelIdRef.current (skip first run)
    MplSlot->>MplComm: setSendHandler(newSendFn)
    MplComm-->>MplSlot: send handler updated
    MplSlot->>MplComm: trigger fakeWs.onopen() again
    MplComm->>Model: new handshake with current model
    Note over MplSlot: Figure DOM NOT rebuilt - toolbar state preserved

    alt Two cells sharing same figure
        Client->>MplSlot: second cell mounts with same figure
        Backend->>Registry: acquire(fig, factory)
        Registry-->>Backend: same cached manager (refcount=2)
        Backend->>Mgr: add_web_socket(second_sync_ws)
        Mgr-->>Backend: both websockets registered
    end

    Note over Client,Fig: CLEANUP / DISPOSE
    
    Backend->>Registry: release(fig) - decrement refcount
    alt Refcount > 0 (other consumers remain)
        Backend->>Mgr: remove_web_socket(sync_ws) only
        Registry-->>Backend: manager kept alive
    else Refcount = 0 (last consumer)
        Registry->>Registry: pop from _managers, _refcounts, _original_geometry
        Registry-->>Backend: return True (destroy)
        Backend->>Mgr: cleanup toolbar callbacks
        Backend->>Fig: restore original dpi & size_inches
        Backend->>Mgr: destroy manager
    end

Loading

_{Re-trigger cubic}

[Copilot]

Pull request overview

This PR improves matplotlib interactive rendering by reusing a single FigureManagerWebAgg per figure across cell reruns (backend) and by rebinding the frontend to new model IDs without tearing down/rebuilding the canvas DOM.

Changes:

• Backend: introduce a weak, refcounted registry to cache/reuse the WebAgg figure manager per figure; destroy it only when the last consumer disposes.
• Frontend: split “mount/build figure once” from “rebind to new model” so reruns don’t clear/recreate the DOM and socket wiring.
• Tests: extend coverage to validate manager reuse/refcounting and frontend rerun rebinding behavior.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
marimo/_plugins/ui/_impl/from_mpl_interactive.py	Adds shared manager registry + refcounted teardown; adjusts cleanup lifecycle to release/destroy managers correctly.
tests/_plugins/ui/_impl/test_from_mpl_interactive.py	Updates cleanup expectations and adds new tests for manager reuse, destruction, and multi-consumer behavior.
frontend/src/plugins/impl/mpl-interactive/MplInteractivePlugin.tsx	Keeps mpl.js figure/socket DOM stable across reruns; adds rebinding logic for new model IDs.
frontend/src/plugins/impl/mpl-interactive/mpl-websocket-shim.ts	Adds setSendHandler to retarget outbound messages without recreating the socket.
frontend/src/plugins/impl/mpl-interactive/__tests__/MplInteractivePlugin.test.tsx	Adds a rerender test ensuring model rebinding doesn’t reconstruct the figure DOM.