added onRemove method
+ small fix