Updates based on PR comments