fix race

Author: qmuntal

go/internal/e2e/mcp_and_agents_test.go

@@ -37,18 +37,13 @@ func TestMCPServers(t *testing.T) {
 		}
 
 		// Simple interaction to verify session works
-		_, err = session.Send(t.Context(), copilot.MessageOptions{
+		message, err := session.SendAndWait(t.Context(), copilot.MessageOptions{
 			Prompt: "What is 2+2?",
 		})
 		if err != nil {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
-		message, err := testharness.GetFinalAssistantMessage(t.Context(), session)
-		if err != nil {
-			t.Fatalf("Failed to get final message: %v", err)
-		}
-
 		if message.Data.Content == nil || !strings.Contains(*message.Data.Content, "4") {
 			t.Errorf("Expected message to contain '4', got: %v", message.Data.Content)
 		}
@@ -168,18 +163,13 @@ func TestCustomAgents(t *testing.T) {
 		}
 
 		// Simple interaction to verify session works
-		_, err = session.Send(t.Context(), copilot.MessageOptions{
+		message, err := session.SendAndWait(t.Context(), copilot.MessageOptions{
 			Prompt: "What is 5+5?",
 		})
 		if err != nil {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
-		message, err := testharness.GetFinalAssistantMessage(t.Context(), session)
-		if err != nil {
-			t.Fatalf("Failed to get final message: %v", err)
-		}
-
 		if message.Data.Content == nil || !strings.Contains(*message.Data.Content, "10") {
 			t.Errorf("Expected message to contain '10', got: %v", message.Data.Content)
 		}
@@ -373,18 +363,13 @@ func TestCombinedConfiguration(t *testing.T) {
 			t.Error("Expected non-empty session ID")
 		}
 
-		_, err = session.Send(t.Context(), copilot.MessageOptions{
+		message, err := session.SendAndWait(t.Context(), copilot.MessageOptions{
 			Prompt: "What is 7+7?",
 		})
 		if err != nil {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
-		message, err := testharness.GetFinalAssistantMessage(t.Context(), session)
-		if err != nil {
-			t.Fatalf("Failed to get final message: %v", err)
-		}
-
 		if message.Data.Content == nil || !strings.Contains(*message.Data.Content, "14") {
 			t.Errorf("Expected message to contain '14', got: %v", message.Data.Content)
 		}