Improved add / remove functionality