Refactor - simplify adding a callback