revert test changes

Author: qmuntal

go/internal/e2e/compaction_test.go

@@ -83,9 +83,6 @@ func TestCompaction(t *testing.T) {
 		if err != nil {
 			t.Fatalf("Failed to send verification message: %v", err)
 		}
-		if answer == nil {
-			t.Fatalf("Expected an answer, got nil")
-		}
 		if answer.Data.Content == nil || !strings.Contains(strings.ToLower(*answer.Data.Content), "dragon") {
 			t.Errorf("Expected answer to contain 'dragon', got %v", answer.Data.Content)
 		}

go/internal/e2e/mcp_and_agents_test.go

@@ -37,15 +37,16 @@ func TestMCPServers(t *testing.T) {
 		}
 
 		// Simple interaction to verify session works
-		message, err := session.SendAndWait(t.Context(), copilot.MessageOptions{
+		_, err = session.Send(t.Context(), copilot.MessageOptions{
 			Prompt: "What is 2+2?",
 		})
 		if err != nil {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
-		if message == nil {
-			t.Fatal("Expected a message, got nil")
+		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") {
@@ -96,10 +97,6 @@ func TestMCPServers(t *testing.T) {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
-		if message == nil {
-			t.Fatalf("Expected a message, got nil")
-		}
-
 		if message.Data.Content == nil || !strings.Contains(*message.Data.Content, "6") {
 			t.Errorf("Expected message to contain '6', got: %v", message.Data.Content)
 		}
@@ -171,15 +168,16 @@ func TestCustomAgents(t *testing.T) {
 		}
 
 		// Simple interaction to verify session works
-		message, err := session.SendAndWait(t.Context(), copilot.MessageOptions{
+		_, err = session.Send(t.Context(), copilot.MessageOptions{
 			Prompt: "What is 5+5?",
 		})
 		if err != nil {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
-		if message == nil {
-			t.Fatal("Expected a message, got nil")
+		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") {
@@ -230,10 +228,6 @@ func TestCustomAgents(t *testing.T) {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
-		if message == nil {
-			t.Fatalf("Expected a message, got nil")
-		}
-
 		if message.Data.Content == nil || !strings.Contains(*message.Data.Content, "12") {
 			t.Errorf("Expected message to contain '12', got: %v", message.Data.Content)
 		}
@@ -379,15 +373,16 @@ func TestCombinedConfiguration(t *testing.T) {
 			t.Error("Expected non-empty session ID")
 		}
 
-		message, err := session.SendAndWait(t.Context(), copilot.MessageOptions{
+		_, err = session.Send(t.Context(), copilot.MessageOptions{
 			Prompt: "What is 7+7?",
 		})
 		if err != nil {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
-		if message == nil {
-			t.Fatalf("Expected a message, got nil")
+		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") {

go/internal/e2e/permissions_test.go

@@ -134,13 +134,18 @@ func TestPermissions(t *testing.T) {
 			t.Fatalf("Failed to write test file: %v", err)
 		}
 
-		_, err = session.SendAndWait(t.Context(), copilot.MessageOptions{
+		_, err = session.Send(t.Context(), copilot.MessageOptions{
 			Prompt: "Edit protected.txt and replace 'protected' with 'hacked'.",
 		})
 		if err != nil {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
+		_, err = testharness.GetFinalAssistantMessage(t.Context(), session)
+		if err != nil {
+			t.Fatalf("Failed to get final message: %v", err)
+		}
+
 		// Verify the file was NOT modified
 		content, err := os.ReadFile(testFile)
 		if err != nil {
@@ -160,13 +165,14 @@ func TestPermissions(t *testing.T) {
 			t.Fatalf("Failed to create session: %v", err)
 		}
 
-		message, err := session.SendAndWait(t.Context(), copilot.MessageOptions{Prompt: "What is 2+2?"})
+		_, err = session.Send(t.Context(), copilot.MessageOptions{Prompt: "What is 2+2?"})
 		if err != nil {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
-		if message == nil {
-			t.Fatal("Expected a message, got nil")
+		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") {