Refactor token handling