Address code review feedback