Refactor session control