Refactor session code