Refactor session logic