Fix token update bug