Unexpected behavior in Matplotlib pane

Before I get started, I am totally excited about the Matplotlib panes supporting ipympl and interactions. However, while I was testing out this new feature, I encountered some unexpected behavior. Here’s an MWE from my testing

import numpy as np
import matplotlib.pyplot as plt
from IPython.display import display

import panel as pn
pn.extension()

TWOPI = 2*np.pi

fig, ax = plt.subplots()
t = np.arange(0.0, TWOPI, 0.001)
initial_amp = 0.5
s = initial_amp*np.sin(t)
l, = ax.plot(t, s, lw=2)
ax.set_xlim(0, TWOPI)
ax.set_ylim(-1, 1)

def add_lines(event):
    xloc = event.xdata
    yloc = event.ydata
    ax.axvline(xloc)
    ax.axhline(yloc)
    fig.canvas.draw_idle()
fig.canvas.mpl_connect('button_press_event', add_lines)

mpl_pane = pn.pane.Matplotlib(fig, dpi=144, interactive=True)
display(mpl_pane)

If I just use plt.show() outside of Panel, I get the interactivity that I expect. However, when I put the figure in a Matplotlib pane, I lose the button_press_event interactivity.

I did find, though, that if I move th mpl_connect statement to after the display function is called, the button_press_event callback is triggered in the Matplotlib pane.

Not sure if this is a bug or a limitation, but it would sure be nice if any event handling in the figure can carry over into the Matplotlib pane without tweaks to pre-existing code.

Thank you for all your hard work on this project!