Fix broken features after previous commit