test(remoteagent): de-flake TestA2ACleanupPropagation

Under CPU contention / low parallelism the test hit two problems:

- The detached CancelTask goroutine was never joined (cancelResultChan was
  created but never read), so a slow RPC could still be in flight at teardown,
  fail on the cancelled t.Context(), and call t.Errorf on a finished test,
  which panics. Both detached goroutines are now joined via a WaitGroup and
  t.Cleanup.

- The parent was cancelled right after the first status event, before the
  remote agent had received any subagent event. In that window the parent does
  not know the subagent's server-assigned task ID, so cancellation cannot
  propagate and subagent cleanup never fires. The test now waits until the
  subagent's own output round-trips to the client before cancelling. The single
  fixed 5s deadline shared across all cleanup waits is also replaced with a
  per-wait, contention-tolerant bound.

Test-only change.

Author: wolo-lab

agent/remoteagent/v2/a2a_e2e_test.go

@@ -26,6 +26,7 @@ import (
 	"net/http"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -380,6 +381,8 @@ func TestA2ACleanupPropagation(t *testing.T) {
 	// Remote A2A server publishes a submitted task and start generating artifact updates
 	// until it detects a context cancelation
 	remoteTaskIDChan, remoteCleanupCalledChan := make(chan a2a.TaskID, 1), make(chan struct{}, 2)
+	// Artifact text the mock subagent streams; the cancel step keys off it.
+	const remoteArtifactText = "remote-subagent-working"
 	serverB := startA2AServer(&mockA2AExecutor{
 		cancelFn: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] {
 			return func(yield func(a2a.Event, error) bool) {
@@ -393,7 +396,7 @@ func TestA2ACleanupPropagation(t *testing.T) {
 					return
 				}
 				for ctx.Err() == nil {
-					if !yield(a2a.NewArtifactEvent(reqCtx, a2a.NewTextPart("foo")), nil) {
+					if !yield(a2a.NewArtifactEvent(reqCtx, a2a.NewTextPart(remoteArtifactText)), nil) {
 						return
 					}
 					time.Sleep(1 * time.Millisecond)
@@ -427,27 +430,44 @@ func TestA2ACleanupPropagation(t *testing.T) {
 
 	client := newA2AClient(t, serverA)
 
-	// Send a streaming message in a detached goroutine, passing status update through chan
+	// Join the detached streaming/cancel goroutines before teardown; a late
+	// t.Errorf from an in-flight RPC on a finished test would panic.
+	var wg sync.WaitGroup
+	t.Cleanup(wg.Wait)
+
+	// remoteStreamingChan closes when the subagent's own output first reaches the client.
 	statusUpdateEventChan := make(chan a2a.Event, 10)
+	remoteStreamingChan := make(chan struct{})
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		defer close(statusUpdateEventChan)
+		remoteStreaming := false
 		msg := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("work"))
 		for event, err := range client.SendStreamingMessage(t.Context(), &a2a.SendMessageRequest{Message: msg}) {
 			if err != nil {
 				t.Errorf("client.SendStreamingMessage() error = %v", err)
 				return
 			}
-			if _, ok := event.(*a2a.TaskArtifactUpdateEvent); ok {
+			if tau, ok := event.(*a2a.TaskArtifactUpdateEvent); ok {
+				if !remoteStreaming && artifactContainsText(tau, remoteArtifactText) {
+					remoteStreaming = true
+					close(remoteStreamingChan)
+				}
 				continue
 			}
 			statusUpdateEventChan <- event
 		}
 	}()
 
-	// Issue a task cancellation request
+	// Cancel only after the subagent's output reaches the client: before that the
+	// parent doesn't know the subagent task ID, so cancellation can't propagate.
 	taskID := (<-statusUpdateEventChan).TaskInfo().TaskID
+	awaitN(t, remoteStreamingChan, 1, "remote subagent streaming")
 	cancelResultChan := make(chan *a2a.Task, 1)
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		defer close(cancelResultChan)
 		task, err := client.CancelTask(t.Context(), &a2a.CancelTaskRequest{ID: taskID})
 		if err != nil {
@@ -470,25 +490,12 @@ func TestA2ACleanupPropagation(t *testing.T) {
 		t.Fatalf("type(lastStreamingUpdate) = %T, want *a2a.TaskStatusUpdateEvent", lastStreamingUpdate)
 	}
 
-	// Check subagent task got cancelled when the parent task was cancelled.
-	// Reads from channel twice because cleanup gets called both for cancelation and execution.
-	timeout := time.After(5 * time.Second)
-	for range 2 {
-		select {
-		case <-remoteCleanupCalledChan:
-		case <-timeout:
-			t.Fatalf("remote cleanup was not called")
-		}
-	}
+	// Subagent cleanup fires twice: once for cancelation, once for execution.
+	// A generous per-wait deadline avoids flaking under CPU contention.
+	awaitN(t, remoteCleanupCalledChan, 2, "remote cleanup")
 	remoteTaskID := <-remoteTaskIDChan
+	awaitN(t, executorCleanupCalledChan, 2, "executor cleanup")
 
-	for range 2 {
-		select {
-		case <-executorCleanupCalledChan:
-		case <-timeout:
-			t.Fatalf("executor cleanup was not called")
-		}
-	}
 	remoteClient := newA2AClient(t, serverB)
 	remoteTask, err := remoteClient.GetTask(t.Context(), &a2a.GetTaskRequest{ID: remoteTaskID})
 	if err != nil {
@@ -497,6 +504,35 @@ func TestA2ACleanupPropagation(t *testing.T) {
 	if remoteTask.Status.State != a2a.TaskStateCanceled {
 		t.Errorf("remoteTask.Status.State = %q, want %q", remoteTask.Status.State, a2a.TaskStateCanceled)
 	}
+
+	// Join the cancel RPC so it can't log on t after the test returns.
+	awaitN(t, cancelResultChan, 1, "cancel task")
+}
+
+// awaitN receives n values from ch or fails the test after a generous, contention-
+// tolerant deadline. A closed channel counts as a receive, so it also joins a
+// goroutine that closed ch without sending.
+func awaitN[T any](t *testing.T, ch <-chan T, n int, what string) {
+	t.Helper()
+	const deadline = 30 * time.Second
+	timer := time.NewTimer(deadline)
+	defer timer.Stop()
+	for i := range n {
+		select {
+		case <-ch:
+		case <-timer.C:
+			t.Fatalf("%s: got %d of %d within %v", what, i, n, deadline)
+		}
+	}
+}
+
+func artifactContainsText(tau *a2a.TaskArtifactUpdateEvent, substr string) bool {
+	for _, p := range tau.Artifact.Parts {
+		if strings.Contains(p.Text(), substr) {
+			return true
+		}
+	}
+	return false
 }
 
 func TestA2ASingleHopFinalResponse(t *testing.T) {