Improvements to handle custom roles