fix up uses of wait_for_agent

- mostly just account for the fact that it returns bool

- also slight refactor to reduce indentation level for readability