Fix edit mode edge case