Fixed some bugs in new feature