fix(adopt): handle CTRL+C cleanly during av adopt (#692)

cmd/av/adopt.go

@@ -70,14 +70,16 @@ The command will adopt the stack of pull requests starting from the specified br
 			if err != nil {
 				return err
 			}
-			return uiutils.RunBubbleTea(&remoteAdoptViewModel{
+			return uiutils.RunBubbleTeaWithContext(ctx, &remoteAdoptViewModel{
+				ctx:        func() context.Context { return ctx },
 				repo:       repo,
 				db:         db,
 				ghClient:   client,
 				branchName: adoptFlags.RemoteBranchName,
 			})
 		}
-		return uiutils.RunBubbleTea(&adoptViewModel{
+		return uiutils.RunBubbleTeaWithContext(ctx, &adoptViewModel{
+			ctx:               func() context.Context { return ctx },
 			repo:              repo,
 			db:                db,
 			currentHEADBranch: plumbing.NewBranchReferenceName(currentBranch),
@@ -143,6 +145,7 @@ func adoptForceAdoption(
 }
 
 type adoptViewModel struct {
+	ctx               func() context.Context
 	repo              *git.Repo
 	db                meta.DB
 	currentHEADBranch plumbing.ReferenceName
@@ -162,6 +165,7 @@ func (vm *adoptViewModel) Init() tea.Cmd {
 func (vm *adoptViewModel) initModel() tea.Cmd {
 	return vm.AddView(
 		actions.NewFindAdoptableLocalBranchesModel(
+			vm.ctx(),
 			vm.repo,
 			vm.db,
 			vm.initTreeSelector,
@@ -247,6 +251,7 @@ func (vm *adoptViewModel) initAdoption(chosenTargets []plumbing.ReferenceName) t
 	}
 	return vm.AddView(
 		actions.NewAdoptBranchesModel(
+			vm.ctx(),
 			vm.db,
 			branches,
 			func() tea.Cmd { return tea.Quit },
@@ -262,6 +267,7 @@ func (vm *adoptViewModel) ExitError() error {
 }
 
 type remoteAdoptViewModel struct {
+	ctx        func() context.Context
 	repo       *git.Repo
 	db         meta.DB
 	ghClient   *gh.Client
@@ -281,6 +287,7 @@ func (vm *remoteAdoptViewModel) Init() tea.Cmd {
 func (vm *remoteAdoptViewModel) initModel() tea.Cmd {
 	return vm.AddView(
 		actions.NewGetRemoteStackedPRModel(
+			vm.ctx(),
 			vm.db.ReadTx().Repository(),
 			vm.ghClient,
 			vm.branchName,
@@ -369,7 +376,7 @@ func (vm *remoteAdoptViewModel) initGitFetch(prs []actions.RemotePRInfo, chosenT
 		refspecs = append(refspecs, fmt.Sprintf("refs/heads/%s:refs/remotes/%s/%s", target.Short(), remote, target.Short()))
 	}
 	return vm.AddView(
-		actions.NewGitFetchModel(vm.repo, refspecs, func() tea.Cmd {
+		actions.NewGitFetchModel(vm.ctx(), vm.repo, refspecs, func() tea.Cmd {
 			return vm.initAdoption(prs, chosenTargets)
 		}),
 	)
@@ -405,6 +412,7 @@ func (vm *remoteAdoptViewModel) initAdoption(prs []actions.RemotePRInfo, chosenT
 	}
 	return vm.AddView(
 		actions.NewAdoptBranchesModel(
+			vm.ctx(),
 			vm.db,
 			branches,
 			func() tea.Cmd {
@@ -417,7 +425,7 @@ func (vm *remoteAdoptViewModel) initAdoption(prs []actions.RemotePRInfo, chosenT
 				for _, ab := range branches {
 					if !hasChild[ab.Name] {
 						// Check out the leaf branch. This can fail if the workspace is dirty. In that case, just quietly exit.
-						_, _ = vm.repo.Git(context.Background(), "switch", ab.Name, "--quiet")
+						_, _ = vm.repo.Git(vm.ctx(), "switch", ab.Name, "--quiet")
 						break
 					}
 				}

cmd/av/main.go

@@ -6,10 +6,12 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"os/signal"
 	"path/filepath"
 	"slices"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 
 	"emperror.dev/errors"
@@ -136,10 +138,14 @@ func run() int {
 	// runtime and various packages (e.g., package init functions).
 	startTime := time.Now()
 	colors.SetupBackgroundColorTypeFromEnv()
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
 	defer func() { _ = cachedRepo.Close() }()
-	err := rootCmd.Execute()
+	err := rootCmd.ExecuteContext(ctx)
 	logrus.WithField("duration", time.Since(startTime)).Debug("command exited")
-	checkCliVersion()
+	if ctx.Err() == nil {
+		checkCliVersion()
+	}
 	var exitSilently actions.ErrExitSilently
 	if errors.As(err, &exitSilently) {
 		return exitSilently.ExitCode

e2e_tests/sigint_test.go

@@ -0,0 +1,98 @@
+package e2e_tests
+
+// TestAdoptExitsOnSIGINT lives outside the testscript harness because
+// testscript v1.10.0 has no public API for sending arbitrary signals to a
+// specifically named background process - it only sends os.Interrupt to
+// every backgrounded command at end-of-test, which is not granular enough
+// to verify that `av adopt` returns promptly after a single SIGINT while
+// the merge-base step is still running.
+//
+// Once the project upgrades to a testscript version with a public
+// signal-by-name primitive (or grows a custom command for it in
+// testscript_test.go), this can move into testdata/script as a .txtar.
+
+import (
+	"bytes"
+	"context"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/aviator-co/av/internal/git/gittest"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAdoptExitsOnSIGINT(t *testing.T) {
+	if _, err := exec.LookPath("sh"); err != nil {
+		t.Skip("requires /bin/sh shim for the fake git binary")
+	}
+
+	repo := gittest.NewTempRepo(t)
+
+	repo.Git(t, "checkout", "-b", "stack-1")
+	repo.CommitFile(t, "my-file", "1a\n", gittest.WithMessage("Commit 1a"))
+	repo.Git(t, "switch", "main")
+
+	realGit, err := exec.LookPath("git")
+	require.NoError(t, err)
+
+	fakeGitDir := t.TempDir()
+	markerPath := filepath.Join(t.TempDir(), "merge-base-started")
+	fakeGitPath := filepath.Join(fakeGitDir, "git")
+	fakeGit := `#!/bin/sh
+if [ "$1" = "merge-base" ]; then
+	if [ -n "$AV_TEST_GIT_MERGE_BASE_STARTED" ]; then
+		: > "$AV_TEST_GIT_MERGE_BASE_STARTED"
+	fi
+	exec sleep 30
+fi
+exec "$REAL_GIT" "$@"
+`
+	require.NoError(t, os.WriteFile(fakeGitPath, []byte(fakeGit), 0o755))
+
+	ctx, cancel := context.WithCancel(t.Context())
+	defer cancel()
+	cmd := exec.CommandContext(ctx, avCmdPath, "--debug", "adopt", "--dry-run")
+	cmd.Dir = repo.RepoDir
+	cmd.Env = append(
+		os.Environ(),
+		"PATH="+fakeGitDir+string(os.PathListSeparator)+os.Getenv("PATH"),
+		"REAL_GIT="+realGit,
+		"AV_TEST_GIT_MERGE_BASE_STARTED="+markerPath,
+	)
+	var stdout, stderr bytes.Buffer
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+
+	require.NoError(t, cmd.Start())
+	waitForFile(t, markerPath, 2*time.Second)
+
+	require.NoError(t, cmd.Process.Signal(os.Interrupt))
+	done := make(chan error, 1)
+	go func() {
+		done <- cmd.Wait()
+	}()
+
+	select {
+	case <-done:
+		// av returned within the deadline; that's the success path. The exit
+		// code is non-zero because the run was cancelled, which is fine.
+	case <-time.After(5 * time.Second):
+		_ = cmd.Process.Kill()
+		t.Fatalf("av adopt did not exit within 5s of SIGINT\nstdout:\n%s\nstderr:\n%s", stdout.String(), stderr.String())
+	}
+}
+
+func waitForFile(t *testing.T, path string, timeout time.Duration) {
+	t.Helper()
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		if _, err := os.Stat(path); err == nil {
+			return
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
+	t.Fatalf("file %s did not appear within %s", path, timeout)
+}

internal/actions/adopt_branches.go

@@ -1,6 +1,7 @@
 package actions
 
 import (
+	"context"
 	"strings"
 
 	"charm.land/bubbles/v2/spinner"
@@ -17,11 +18,13 @@ type AdoptingBranch struct {
 }
 
 func NewAdoptBranchesModel(
+	ctx context.Context,
 	db meta.DB,
 	branches []AdoptingBranch,
 	onDone func() tea.Cmd,
 ) tea.Model {
 	return &AdoptBranchesModel{
+		ctx:      func() context.Context { return ctx },
 		db:       db,
 		spinner:  spinner.New(spinner.WithSpinner(spinner.Dot)),
 		branches: branches,
@@ -30,6 +33,7 @@ func NewAdoptBranchesModel(
 }
 
 type AdoptBranchesModel struct {
+	ctx      func() context.Context
 	db       meta.DB
 	spinner  spinner.Model
 	branches []AdoptingBranch
@@ -68,6 +72,9 @@ func (m *AdoptBranchesModel) View() tea.View {
 func (m *AdoptBranchesModel) adoptBranches() tea.Msg {
 	tx := m.db.WriteTx()
 	for _, branch := range m.branches {
+		if err := m.ctx().Err(); err != nil {
+			return err
+		}
 		bi, _ := tx.Branch(branch.Name)
 		bi.Parent = branch.Parent
 		bi.PullRequest = branch.PullRequest
@@ -76,6 +83,9 @@ func (m *AdoptBranchesModel) adoptBranches() tea.Msg {
 		}
 		tx.SetBranch(bi)
 	}
+	if err := m.ctx().Err(); err != nil {
+		return err
+	}
 	if err := tx.Commit(); err != nil {
 		return err
 	}

internal/actions/find_adoptable_local_branches.go

@@ -15,6 +15,7 @@ import (
 )
 
 func NewFindAdoptableLocalBranchesModel(
+	ctx context.Context,
 	repo *git.Repo,
 	db meta.DB,
 	onDone func(
@@ -24,6 +25,7 @@ func NewFindAdoptableLocalBranchesModel(
 	) tea.Cmd,
 ) tea.Model {
 	return &FindAdoptableLocalBranchesModel{
+		ctx:     func() context.Context { return ctx },
 		repo:    repo,
 		db:      db,
 		spinner: spinner.New(spinner.WithSpinner(spinner.Dot)),
@@ -32,6 +34,7 @@ func NewFindAdoptableLocalBranchesModel(
 }
 
 type FindAdoptableLocalBranchesModel struct {
+	ctx     func() context.Context
 	repo    *git.Repo
 	db      meta.DB
 	spinner spinner.Model
@@ -49,7 +52,7 @@ func (m *FindAdoptableLocalBranchesModel) Init() tea.Cmd {
 		if err != nil {
 			return err
 		}
-		branches, err := treedetector.DetectBranches(context.Background(), m.repo, unmanagedBranches)
+		branches, err := treedetector.DetectBranches(m.ctx(), m.repo, unmanagedBranches)
 		if err != nil {
 			return err
 		}