🚦 Improve Workflow Polling and Signal Handling (#4002)

* fixes to waiting and repo reuse

* update disable logic and messages

* Update pkg/cli/signal_aware_poll.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: dsyme

pkg/cli/enable.go

@@ -250,3 +250,105 @@ func toggleWorkflowsByNames(workflowNames []string, enable bool) error {
 
 	return nil
 }
+
+// DisableAllWorkflowsExcept disables all workflows except the specified ones
+// Typically used to disable all workflows except the one being trialled
+func DisableAllWorkflowsExcept(repoSlug string, exceptWorkflows []string, verbose bool) error {
+	workflowsDir := ".github/workflows"
+
+	// Check if workflows directory exists
+	if _, err := os.Stat(workflowsDir); os.IsNotExist(err) {
+		if verbose {
+			fmt.Fprintln(os.Stderr, console.FormatInfoMessage("No .github/workflows directory found, nothing to disable"))
+		}
+		return nil
+	}
+
+	// Get all .yml and .yaml files
+	ymlFiles, _ := filepath.Glob(filepath.Join(workflowsDir, "*.yml"))
+	yamlFiles, _ := filepath.Glob(filepath.Join(workflowsDir, "*.yaml"))
+	allYAMLFiles := append(ymlFiles, yamlFiles...)
+
+	if len(allYAMLFiles) == 0 {
+		if verbose {
+			fmt.Fprintln(os.Stderr, console.FormatInfoMessage("No YAML workflow files found"))
+		}
+		return nil
+	}
+
+	// Create a set of workflows to keep enabled
+	keepEnabled := make(map[string]bool)
+	for _, workflowName := range exceptWorkflows {
+		// Add both .md and .lock.yml variants
+		keepEnabled[workflowName+".md"] = true
+		keepEnabled[workflowName+".lock.yml"] = true
+		keepEnabled[workflowName] = true // In case the full filename is provided
+	}
+
+	// Filter to find workflows to disable
+	var workflowsToDisable []string
+
+	for _, yamlFile := range allYAMLFiles {
+		base := filepath.Base(yamlFile)
+
+		// Skip if it's in the keep-enabled set
+		if keepEnabled[base] {
+			if verbose {
+				fmt.Fprintf(os.Stderr, "Keeping enabled: %s\n", base)
+			}
+			continue
+		}
+
+		// Check if the base name without extension matches
+		nameWithoutExt := strings.TrimSuffix(base, filepath.Ext(base))
+		if keepEnabled[nameWithoutExt] {
+			if verbose {
+				fmt.Fprintf(os.Stderr, "Keeping enabled: %s\n", base)
+			}
+			continue
+		}
+
+		workflowsToDisable = append(workflowsToDisable, base)
+	}
+
+	if len(workflowsToDisable) == 0 {
+		if verbose {
+			fmt.Fprintln(os.Stderr, console.FormatInfoMessage("No workflows to disable"))
+		}
+		return nil
+	}
+
+	// Show what will be disabled
+	fmt.Fprintf(os.Stderr, "Disabling %d workflow(s) in cloned repository:\n", len(workflowsToDisable))
+	for _, workflow := range workflowsToDisable {
+		fmt.Fprintf(os.Stderr, "  %s\n", workflow)
+	}
+
+	// Disable each workflow
+	var failures []string
+	for _, workflow := range workflowsToDisable {
+		args := []string{"workflow", "disable", workflow}
+		if repoSlug != "" {
+			args = append(args, "--repo", repoSlug)
+		}
+
+		cmd := exec.Command("gh", args...)
+		if output, err := cmd.CombinedOutput(); err != nil {
+			if verbose {
+				fmt.Fprintf(os.Stderr, "Warning: Failed to disable workflow %s: %v\n%s\n", workflow, err, string(output))
+			}
+			failures = append(failures, workflow)
+		} else {
+			if verbose {
+				fmt.Fprintf(os.Stderr, "Disabled workflow: %s\n", workflow)
+			}
+		}
+	}
+
+	if len(failures) > 0 {
+		return fmt.Errorf("failed to disable %d workflow(s): %s", len(failures), strings.Join(failures, ", "))
+	}
+
+	fmt.Fprintln(os.Stderr, console.FormatSuccessMessage(fmt.Sprintf("Disabled %d workflow(s)", len(workflowsToDisable))))
+	return nil
+}

pkg/cli/pr_automerge.go

@@ -120,52 +120,41 @@ func AutoMergePullRequestsLegacy(repoSlug string, verbose bool) error {
 func WaitForWorkflowCompletion(repoSlug, runID string, timeoutMinutes int, verbose bool) error {
 	prAutomergeLog.Printf("Waiting for workflow completion: repo=%s, runID=%s, timeout=%d minutes", repoSlug, runID, timeoutMinutes)
 
-	if verbose {
-		fmt.Fprintln(os.Stderr, console.FormatInfoMessage(fmt.Sprintf("Waiting for workflow completion (timeout: %d minutes)", timeoutMinutes)))
-	}
-
-	// Use the repository slug directly
-	fullRepoName := repoSlug
-
 	timeout := time.Duration(timeoutMinutes) * time.Minute
-	start := time.Now()
 
-	for {
-		// Check if timeout exceeded
-		if time.Since(start) > timeout {
-			return fmt.Errorf("workflow execution timed out after %d minutes", timeoutMinutes)
-		}
-
-		// Check workflow status
-		cmd := exec.Command("gh", "run", "view", runID, "--repo", fullRepoName, "--json", "status,conclusion")
-		output, err := cmd.Output()
+	return PollWithSignalHandling(PollOptions{
+		PollInterval: 10 * time.Second,
+		Timeout:      timeout,
+		PollFunc: func() (PollResult, error) {
+			// Check workflow status
+			cmd := exec.Command("gh", "run", "view", runID, "--repo", repoSlug, "--json", "status,conclusion")
+			output, err := cmd.Output()
 
-		if err != nil {
-			return fmt.Errorf("failed to check workflow status: %w", err)
-		}
-
-		status := string(output)
+			if err != nil {
+				return PollFailure, fmt.Errorf("failed to check workflow status: %w", err)
+			}
 
-		// Check if completed
-		if strings.Contains(status, `"status":"completed"`) {
-			if strings.Contains(status, `"conclusion":"success"`) {
-				if verbose {
-					fmt.Fprintln(os.Stderr, console.FormatSuccessMessage("Workflow completed successfully"))
+			status := string(output)
+
+			// Check if completed
+			if strings.Contains(status, `"status":"completed"`) {
+				if strings.Contains(status, `"conclusion":"success"`) {
+					return PollSuccess, nil
+				} else if strings.Contains(status, `"conclusion":"failure"`) {
+					return PollFailure, fmt.Errorf("workflow failed")
+				} else if strings.Contains(status, `"conclusion":"cancelled"`) {
+					return PollFailure, fmt.Errorf("workflow was cancelled")
+				} else {
+					return PollFailure, fmt.Errorf("workflow completed with unknown conclusion")
 				}
-				return nil
-			} else if strings.Contains(status, `"conclusion":"failure"`) {
-				return fmt.Errorf("workflow failed")
-			} else if strings.Contains(status, `"conclusion":"cancelled"`) {
-				return fmt.Errorf("workflow was cancelled")
-			} else {
-				return fmt.Errorf("workflow completed with unknown conclusion")
 			}
-		}
 
-		// Still running, wait before checking again
-		if verbose {
-			fmt.Fprintln(os.Stderr, console.FormatProgressMessage("Workflow still running..."))
-		}
-		time.Sleep(10 * time.Second)
-	}
+			// Still running, continue polling
+			return PollContinue, nil
+		},
+		StartMessage:    fmt.Sprintf("Waiting for workflow completion (timeout: %d minutes)", timeoutMinutes),
+		ProgressMessage: "Workflow still running...",
+		SuccessMessage:  "Workflow completed successfully",
+		Verbose:         verbose,
+	})
 }

pkg/cli/pr_automerge_test.go

@@ -0,0 +1,23 @@
+package cli
+
+import (
+	"testing"
+)
+
+// TestWaitForWorkflowCompletionUsesSignalHandling verifies that WaitForWorkflowCompletion
+// uses the signal-aware polling helper, which provides Ctrl-C support
+func TestWaitForWorkflowCompletionUsesSignalHandling(t *testing.T) {
+	// This test verifies that the function uses PollWithSignalHandling
+	// by checking that it times out correctly (a key feature of the helper)
+
+	// We can't easily test the actual workflow checking without a real workflow,
+	// but we can verify that the timeout mechanism works, which confirms
+	// it's using the polling helper
+
+	err := WaitForWorkflowCompletion("nonexistent/repo", "12345", 0, false)
+
+	// Should timeout or fail to check workflow status
+	if err == nil {
+		t.Error("Expected error for nonexistent workflow, got nil")
+	}
+}

pkg/cli/signal_aware_poll.go

@@ -0,0 +1,110 @@
+package cli
+
+import (
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"github.com/githubnext/gh-aw/pkg/console"
+	"github.com/githubnext/gh-aw/pkg/logger"
+)
+
+var pollLog = logger.New("cli:signal_aware_poll")
+
+// PollResult represents the result of a polling operation
+type PollResult int
+
+const (
+	// PollContinue indicates polling should continue
+	PollContinue PollResult = iota
+	// PollSuccess indicates polling completed successfully
+	PollSuccess
+	// PollFailure indicates polling failed
+	PollFailure
+)
+
+// PollOptions contains configuration for signal-aware polling
+type PollOptions struct {
+	// Interval between poll attempts
+	PollInterval time.Duration
+	// Timeout for the entire polling operation
+	Timeout time.Duration
+	// Function to call on each poll iteration
+	// Should return PollContinue to keep polling, PollSuccess to succeed, or PollFailure to fail
+	PollFunc func() (PollResult, error)
+	// Message to display when polling starts (optional)
+	StartMessage string
+	// Message to display on each poll iteration (optional)
+	ProgressMessage string
+	// Message to display on successful completion (optional)
+	SuccessMessage string
+	// Whether to show verbose progress messages
+	Verbose bool
+}
+
+// PollWithSignalHandling polls with a function until it succeeds, fails, times out, or receives an interrupt signal
+// This provides a reusable pattern for any operation that needs to poll with graceful Ctrl-C handling
+func PollWithSignalHandling(options PollOptions) error {
+	pollLog.Printf("Starting polling: interval=%v, timeout=%v", options.PollInterval, options.Timeout)
+
+	if options.Verbose && options.StartMessage != "" {
+		fmt.Fprintln(os.Stderr, console.FormatInfoMessage(options.StartMessage))
+	}
+
+	// Set up signal handling for graceful shutdown
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+	defer signal.Stop(sigChan)
+
+	// Set up timeout
+	start := time.Now()
+	ticker := time.NewTicker(options.PollInterval)
+	defer ticker.Stop()
+
+	// Perform initial check immediately
+	result, err := options.PollFunc()
+	if result == PollSuccess {
+		if options.Verbose && options.SuccessMessage != "" {
+			fmt.Fprintln(os.Stderr, console.FormatSuccessMessage(options.SuccessMessage))
+		}
+		return nil
+	} else if result == PollFailure {
+		return err
+	}
+
+	// Continue polling
+	for {
+		select {
+		case <-sigChan:
+			pollLog.Print("Received interrupt signal")
+			fmt.Fprintln(os.Stderr, console.FormatInfoMessage("Received interrupt signal, stopping wait..."))
+			return fmt.Errorf("interrupted by user")
+
+		case <-ticker.C:
+			// Check if timeout exceeded
+			if options.Timeout > 0 && time.Since(start) > options.Timeout {
+				pollLog.Printf("Timeout exceeded: %v", options.Timeout)
+				return fmt.Errorf("operation timed out after %v", options.Timeout)
+			}
+
+			// Poll for status
+			result, err := options.PollFunc()
+
+			if result == PollSuccess {
+				if options.Verbose && options.SuccessMessage != "" {
+					fmt.Fprintln(os.Stderr, console.FormatSuccessMessage(options.SuccessMessage))
+				}
+				return nil
+			} else if result == PollFailure {
+				return err
+			}
+
+			// Still waiting, show progress if enabled
+			if options.Verbose && options.ProgressMessage != "" {
+				fmt.Fprintln(os.Stderr, console.FormatProgressMessage(options.ProgressMessage))
+			}
+		}
+	}
+}