More improvements to edit modal button handling