agent-tools/cmd/agentws/prune.go

package main

import (
	"errors"
	"fmt"
	"io"
	"path/filepath"
	"sort"

	"git.unkin.net/unkin/agent-tools/internal/agent"

	"github.com/spf13/cobra"
)

// Verdicts a worktree can be classified into.
const (
	verdictKeep         = "keep"
	verdictRemove       = "remove"
	verdictRemoveBranch = "remove+branch"
)

// prLister is the slice of the Gitea client prune needs, so tests can drive
// classification without a live server.
type prLister interface {
	ListPRs(repoPath, state string) ([]agent.PullRequest, error)
}

type pruneResult struct {
	wt      managedWt
	verdict string
	reason  string
	// proven records that git itself confirmed the branch's commits survive
	// elsewhere; only then may a branch delete override git's own guard.
	proven bool
}

// repoCtx is the per-repo state classification is decided against.
type repoCtx struct {
	srcDir    string
	defBranch string
	prs       map[string]agent.PullRequest
	prsKnown  bool
	// fetched records that this run's pruning fetch succeeded; without it an
	// origin/<branch> ref may be stale and due for deletion, so it proves nothing.
	fetched bool
}

func newPruneCmd() *cobra.Command {
	var apply, keepBranches bool
	cmd := &cobra.Command{
		Use:          "prune",
		Short:        "Classify managed worktrees and remove the ones whose work is safely upstream",
		Long:         "prune inspects every managed worktree, classifies it against git and its Gitea\npull request, and removes the ones whose work is provably upstream. It is a dry\nrun unless --yes is given.",
		SilenceUsage: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runPrune(cmd.OutOrStdout(), pruneClient(), apply, keepBranches)
		},
	}
	f := cmd.Flags()
	f.BoolVar(&apply, "yes", false, "Actually remove worktrees (default is a dry run)")
	f.BoolVar(&keepBranches, "keep-branches", false, "Never delete a local branch, whatever the classification")
	return cmd
}

// pruneClient builds a Gitea client, falling back to anonymous access when no
// token can be minted; prune degrades to git-only signals if that fails too.
func pruneClient() prLister {
	tok, err := agent.GiteaToken()
	if err != nil {
		tok = ""
	}
	return agent.NewGiteaClient(tok)
}

func runPrune(out io.Writer, prs prLister, apply, keepBranches bool) error {
	managed, err := managedWorktrees()
	if err != nil {
		return err
	}
	if len(managed) == 0 {
		_, _ = fmt.Fprintln(out, "no managed worktrees")
		return nil
	}

	byRepo := map[string][]managedWt{}
	for _, w := range managed {
		byRepo[w.srcDir] = append(byRepo[w.srcDir], w)
	}
	srcDirs := make([]string, 0, len(byRepo))
	for dir := range byRepo {
		srcDirs = append(srcDirs, dir)
	}
	sort.Strings(srcDirs)

	var results []pruneResult
	for _, srcDir := range srcDirs {
		ctx, err := newRepoCtx(out, prs, srcDir)
		if err != nil {
			for _, w := range byRepo[srcDir] {
				results = append(results, pruneResult{wt: w, verdict: verdictKeep, reason: "repo state unknown: " + err.Error()})
			}
			continue
		}
		for _, w := range byRepo[srcDir] {
			res, err := classify(w, ctx)
			if err != nil {
				res = pruneResult{wt: w, verdict: verdictKeep, reason: "inspection failed: " + err.Error()}
			}
			results = append(results, res)
		}
	}

	for _, r := range results {
		_, _ = fmt.Fprintf(out, "%-44s %-34s %-14s %s\n", filepath.Base(r.wt.path), r.wt.branch, plannedVerdict(r, keepBranches), r.reason)
	}
	if !apply {
		_, _ = fmt.Fprintln(out, "dry run: nothing removed (pass --yes to apply)")
		return nil
	}

	var errs []error
	for _, r := range results {
		if r.verdict == verdictKeep {
			continue
		}
		deleteBranch := r.verdict == verdictRemoveBranch && !keepBranches
		if err := removeWorktree(out, r.wt, deleteBranch, r.proven); err != nil {
			errs = append(errs, fmt.Errorf("%s: %w", r.wt.path, err))
		}
	}
	return errors.Join(errs...)
}

// plannedVerdict is what will actually happen, so --keep-branches does not
// print a branch deletion it will not perform.
func plannedVerdict(r pruneResult, keepBranches bool) string {
	if keepBranches && r.verdict == verdictRemoveBranch {
		return verdictRemove
	}
	return r.verdict
}

// newRepoCtx refreshes a source repo and collects the signals prune classifies
// against. A failed fetch or an unreachable Gitea is reported and tolerated:
// the signals that hold offline still work, and the rest are recorded as
// unverified.
func newRepoCtx(out io.Writer, prs prLister, srcDir string) (repoCtx, error) {
	ctx := repoCtx{srcDir: srcDir, prs: map[string]agent.PullRequest{}}
	repo := filepath.Base(srcDir)
	if err := agent.GitFetchPrune(srcDir, "origin", credentialHelperArgs()...); err != nil {
		_, _ = fmt.Fprintf(out, "warn: fetch %s: %v (remote state unverified)\n", repo, err)
	} else {
		ctx.fetched = true
	}
	def, err := agent.GitRemoteDefaultBranch(srcDir, "origin")
	if err != nil {
		return repoCtx{}, err
	}
	ctx.defBranch = def

	if prs == nil {
		return ctx, nil
	}
	list, err := prs.ListPRs(repoPath(srcDir, repo), "all")
	switch {
	case errors.Is(err, agent.ErrPRListTruncated):
		// A branch missing from a partial listing must not read as "no PR".
		_, _ = fmt.Fprintf(out, "warn: list PRs for %s: %v (older PRs unseen)\n", repo, err)
		ctx.prs = prsByBranch(list)
	case err != nil:
		_, _ = fmt.Fprintf(out, "warn: list PRs for %s: %v (git signals only)\n", repo, err)
	default:
		ctx.prs = prsByBranch(list)
		ctx.prsKnown = true
	}
	return ctx, nil
}

// repoPath is the Gitea "owner/repo" for a checkout, read from origin's URL
// because not every managed repo lives under AGENTWS_OWNER.
func repoPath(srcDir, repo string) string {
	url, err := agent.GitRemoteURL(srcDir, "origin")
	if err == nil && agent.RemoteHost(url) == giteaHost() {
		if path, err := agent.RepoPathFromRemoteURL(url); err == nil {
			return path
		}
	}
	return owner() + "/" + repo
}

// prsByBranch indexes PRs by head branch, preferring an open PR and otherwise
// the most recent one when a branch has been used more than once.
func prsByBranch(list []agent.PullRequest) map[string]agent.PullRequest {
	out := map[string]agent.PullRequest{}
	for _, pr := range list {
		branch := agent.PRHeadBranch(pr)
		if branch == "" {
			continue
		}
		if cur, ok := out[branch]; ok && !supersedes(pr, cur) {
			continue
		}
		out[branch] = pr
	}
	return out
}

func supersedes(a, b agent.PullRequest) bool {
	if a.IsOpen() != b.IsOpen() {
		return a.IsOpen()
	}
	if a.Merged != b.Merged {
		return a.Merged
	}
	return a.Number > b.Number
}

// headContainedIn reports whether the worktree's HEAD is reachable from ref. A
// ref that cannot be resolved proves nothing, so it reads as not contained.
func headContainedIn(dir, ref string) bool {
	if ref == "" {
		return false
	}
	ok, err := agent.GitIsAncestor(dir, "HEAD", ref)
	return err == nil && ok
}

// classify applies the prune precedence: dirty and open-PR worktrees are kept,
// provably-upstream work loses its branch too, and anything unproven keeps its
// branch so no commits become unreachable. A PR's state alone never authorises
// deleting a branch — git must confirm HEAD is contained in what merged or in
// what origin still holds, and origin's refs only count when this run's pruning
// fetch refreshed them.
func classify(wt managedWt, ctx repoCtx) (pruneResult, error) {
	res := pruneResult{wt: wt}
	dirty, err := agent.GitIsDirty(wt.path)
	if err != nil {
		return res, err
	}
	if dirty {
		res.verdict, res.reason = verdictKeep, "dirty"
		return res, nil
	}

	pr, hasPR := ctx.prs[wt.branch]
	if hasPR && pr.IsOpen() {
		res.verdict, res.reason = verdictKeep, fmt.Sprintf("PR open #%d", pr.Number)
		return res, nil
	}

	upstream := "origin/" + ctx.defBranch
	contained, err := agent.GitIsAncestor(wt.path, "HEAD", upstream)
	if err != nil {
		return res, err
	}
	if contained {
		res.verdict, res.reason, res.proven = verdictRemoveBranch, "contained in "+upstream, true
		return res, nil
	}
	unmerged, err := agent.GitUnmergedCommits(wt.path, upstream, "HEAD")
	if err != nil {
		return res, err
	}
	if unmerged == 0 {
		// git cherry proves the patches reached that history, not that they stand at its tip.
		res.verdict, res.reason, res.proven = verdictRemoveBranch, "patch-equivalent commits in "+upstream+" history", true
		return res, nil
	}

	remote := "origin/" + wt.branch
	onOrigin := hasPR && ctx.fetched && agent.GitRemoteBranchExists(ctx.srcDir, "origin", wt.branch)
	switch {
	case hasPR && pr.Merged && headContainedIn(wt.path, pr.Head.Sha):
		res.verdict, res.reason, res.proven = verdictRemoveBranch, fmt.Sprintf("PR merged #%d, HEAD contained in the merged head", pr.Number), true
	case hasPR && pr.Merged && onOrigin && headContainedIn(wt.path, remote):
		res.verdict, res.reason, res.proven = verdictRemoveBranch, fmt.Sprintf("PR merged #%d, HEAD contained in %s", pr.Number, remote), true
	case hasPR && pr.Merged && !ctx.fetched:
		res.verdict, res.reason = verdictRemove, fmt.Sprintf("PR merged #%d, fetch failed so %s is unverified", pr.Number, remote)
	case hasPR && pr.Merged:
		res.verdict, res.reason = verdictRemove, fmt.Sprintf("PR merged #%d, local commits not in the merged head", pr.Number)
	case hasPR && onOrigin && headContainedIn(wt.path, remote):
		res.verdict, res.reason, res.proven = verdictRemoveBranch, fmt.Sprintf("PR closed #%d, HEAD contained in %s", pr.Number, remote), true
	case hasPR && onOrigin:
		res.verdict, res.reason = verdictRemove, fmt.Sprintf("PR closed #%d, local commits not on %s", pr.Number, remote)
	case hasPR && !ctx.fetched:
		res.verdict, res.reason = verdictRemove, fmt.Sprintf("PR closed #%d, fetch failed so %s is unverified", pr.Number, remote)
	case hasPR:
		res.verdict, res.reason = verdictRemove, fmt.Sprintf("PR closed #%d, branch gone", pr.Number)
	case ctx.prsKnown:
		res.verdict, res.reason = verdictRemove, "no PR"
	default:
		res.verdict, res.reason = verdictRemove, "PR state unknown"
	}
	return res, nil
}