Address agent code review comments