[minor] Add custom callback support