Improve session handling (part 2)