Fix code review feedback