Improve context infrastructure