Track agent usage of contexts

**Describe the enhancement:**

The agent uses `context.Background` in a few places. We should check to see if we can pass a context instead or use `context.TODO` in places where we should.

**Describe a specific use case for the enhancement or feature:**
Allowing proper contexts to be passed may allow us to cancel in-progress operations if they exceed a deadline or timeout