diff --git a/loop/agent_test.go b/loop/agent_test.go
index c62fd21..9663e26 100644
--- a/loop/agent_test.go
+++ b/loop/agent_test.go
@@ -94,26 +94,16 @@
 
 	// Collect responses with a timeout
 	var responses []AgentMessage
-	timeout := time.After(10 * time.Second)
+	ctx2, _ := context.WithDeadline(ctx, time.Now().Add(10*time.Second))
 	done := false
+	it := agent.NewIterator(ctx2, 0)
 
 	for !done {
-		select {
-		case <-timeout:
-			t.Log("Timeout reached while waiting for agent responses")
+		msg := it.Next()
+		t.Logf("Received message: Type=%s, EndOfTurn=%v, Content=%q", msg.Type, msg.EndOfTurn, msg.Content)
+		responses = append(responses, *msg)
+		if msg.EndOfTurn {
 			done = true
-		default:
-			select {
-			case msg := <-agent.outbox:
-				t.Logf("Received message: Type=%s, EndOfTurn=%v, Content=%q", msg.Type, msg.EndOfTurn, msg.Content)
-				responses = append(responses, msg)
-				if msg.EndOfTurn {
-					done = true
-				}
-			default:
-				// No more messages available right now
-				time.Sleep(100 * time.Millisecond)
-			}
 		}
 	}