Fix brekage from previous patches