agent-tools/cmd/agentws/prune.go

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"text/tabwriter"

	"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
	// unfetched explains why, so a verdict can say which it was.
	unfetched string
}

// pruneOpts is the knob set runPrune is driven by.
type pruneOpts struct {
	apply            bool
	keepBranches     bool
	noFetch          bool
	jsonOut          bool
	includeKeep      bool
	includeUnmanaged bool
}

// reportEntry is the --json shape: one object per worktree, mirroring the table.
type reportEntry struct {
	Repo    string `json:"repo"`
	Branch  string `json:"branch"`
	Path    string `json:"path"`
	Verdict string `json:"verdict"`
	Reason  string `json:"reason"`
	Managed bool   `json:"managed"`
	Applied bool   `json:"applied"`
}

func newPruneCmd() *cobra.Command {
	var opts pruneOpts
	cmd := &cobra.Command{
		Use:          "prune",
		Short:        "Classify worktrees and remove the ones whose work is safely upstream",
		Long:         "prune inspects every worktree it can find — the managed ones under the worktree\nroot plus whatever `git worktree list` reports for the source checkouts — and\nclassifies each against git and its Gitea pull request. It reports and changes\nnothing unless --yes is given.",
		SilenceUsage: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runPrune(cmd.OutOrStdout(), cmd.ErrOrStderr(), pruneClient(), opts)
		},
	}
	f := cmd.Flags()
	f.BoolVar(&opts.apply, "yes", false, "Actually remove worktrees (default is a dry run)")
	f.BoolVar(&opts.keepBranches, "keep-branches", false, "Never delete a local branch, whatever the classification")
	f.BoolVar(&opts.noFetch, "no-fetch", false, "Do not fetch; judge against the refs already on disk")
	f.BoolVar(&opts.jsonOut, "json", false, "Emit JSON instead of a table")
	f.BoolVar(&opts.includeKeep, "include-keep", false, "Dangerous: also remove worktrees classified keep (needs --yes). Destroys uncommitted changes and paused rebase/merge state, which no branch is carrying; only the branch itself survives")
	f.BoolVar(&opts.includeUnmanaged, "include-unmanaged", false, "Also remove worktrees that live outside the worktree root")
	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, errOut io.Writer, prs prLister, opts pruneOpts) error {
	// In JSON mode stdout carries the document alone, so notes go to stderr.
	notes := out
	if opts.jsonOut {
		notes = errOut
	}

	worktrees, err := allWorktrees()
	if err != nil {
		return err
	}
	if len(worktrees) == 0 {
		if opts.jsonOut {
			_, _ = fmt.Fprintln(out, "[]")
			return nil
		}
		_, _ = fmt.Fprintln(out, "no managed worktrees")
		return nil
	}

	results := classifyAll(notes, prs, worktrees, opts)
	if err := report(out, results, opts); err != nil {
		return err
	}
	if !opts.apply {
		if !opts.jsonOut {
			_, _ = fmt.Fprintln(out, "dry run: nothing removed (pass --yes to apply)")
		}
		return nil
	}
	return applyPrune(notes, results, opts)
}

// classifyAll groups worktrees by source repo so each repo is fetched and its
// PRs listed once, then classifies every worktree against that repo's state.
func classifyAll(notes io.Writer, prs prLister, worktrees []managedWt, opts pruneOpts) []pruneResult {
	byRepo := map[string][]managedWt{}
	var results []pruneResult
	for _, w := range worktrees {
		switch {
		case w.inspectErr != nil:
			// Unknown is not gone: a checkout git refused to answer for keeps.
			results = append(results, pruneResult{wt: w, verdict: verdictKeep, reason: "inspection failed: " + oneLine(w.inspectErr.Error())})
		case w.orphan:
			results = append(results, pruneResult{wt: w, verdict: verdictRemove, reason: "backing repo gone, no git state to read"})
		default:
			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)

	for _, srcDir := range srcDirs {
		ctx, err := newRepoCtx(notes, prs, srcDir, opts.noFetch)
		if err != nil {
			for _, w := range byRepo[srcDir] {
				results = append(results, pruneResult{wt: w, verdict: verdictKeep, reason: "repo state unknown: " + oneLine(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: " + oneLine(err.Error())}
			}
			results = append(results, res)
		}
	}
	sort.SliceStable(results, func(i, j int) bool {
		if results[i].wt.repo != results[j].wt.repo {
			return results[i].wt.repo < results[j].wt.repo
		}
		return results[i].wt.path < results[j].wt.path
	})
	return results
}

// report writes the classification as a table or as JSON.
func report(out io.Writer, results []pruneResult, opts pruneOpts) error {
	if opts.jsonOut {
		entries := make([]reportEntry, 0, len(results))
		for _, r := range results {
			entries = append(entries, reportEntry{
				Repo:    r.wt.repo,
				Branch:  r.wt.branch,
				Path:    r.wt.path,
				Verdict: plannedVerdict(r, opts),
				Reason:  r.reason,
				Managed: r.wt.managed,
				Applied: opts.apply && willRemove(r, opts),
			})
		}
		enc := json.NewEncoder(out)
		enc.SetIndent("", "  ")
		return enc.Encode(entries)
	}

	// Padded with spaces only, so the table reads the same with or without a TTY.
	tw := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
	_, _ = fmt.Fprintln(tw, "REPO\tBRANCH\tPATH\tVERDICT\tREASON")
	for _, r := range results {
		_, _ = fmt.Fprintf(tw, "%s\t%s\t%s\t%s\t%s\n",
			dash(r.wt.repo), dash(r.wt.branch), abbrevHome(r.wt.path), plannedVerdict(r, opts), r.reason)
	}
	return tw.Flush()
}

// oneLine flattens a git error onto a single line so one row stays one row.
func oneLine(s string) string {
	return strings.Join(strings.Fields(s), " ")
}

func dash(s string) string {
	if s == "" {
		return "-"
	}
	return s
}

// abbrevHome shortens $HOME to ~ so paths do not dominate the table.
func abbrevHome(path string) string {
	home, err := os.UserHomeDir()
	if err != nil || home == "" || !strings.HasPrefix(path, home+string(filepath.Separator)) {
		return path
	}
	return "~" + path[len(home):]
}

// applyPrune performs the removals the classification authorised, reporting each
// one and collecting failures so one bad worktree does not stop the rest.
func applyPrune(out io.Writer, results []pruneResult, opts pruneOpts) error {
	var errs []error
	for _, r := range results {
		if !willRemove(r, opts) {
			if !r.wt.managed && r.verdict != verdictKeep {
				_, _ = fmt.Fprintf(out, "skipped %s: outside the worktree root (pass --include-unmanaged)\n", r.wt.path)
			}
			continue
		}
		if r.verdict == verdictKeep {
			_, _ = fmt.Fprintf(out, "warn: removing %s despite %q (--include-keep)\n", r.wt.path, r.reason)
		}
		deleteBranch := plannedVerdict(r, opts) == verdictRemoveBranch
		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...)
}

// willRemove is the single gate on destruction: a keep verdict needs
// --include-keep, and a worktree outside the worktree root needs
// --include-unmanaged.
func willRemove(r pruneResult, opts pruneOpts) bool {
	if !r.wt.managed && !opts.includeUnmanaged {
		return false
	}
	if r.verdict == verdictKeep {
		return opts.includeKeep
	}
	return true
}

// plannedVerdict is what will actually happen, so no flag prints an action it
// will not perform. A detached HEAD has no branch to delete, and a keep forced
// through with --include-keep never takes its branch with it.
func plannedVerdict(r pruneResult, opts pruneOpts) string {
	if r.verdict != verdictRemoveBranch {
		return r.verdict
	}
	if opts.keepBranches || r.wt.detached {
		return verdictRemove
	}
	return verdictRemoveBranch
}

// newRepoCtx refreshes a source repo and collects the signals prune classifies
// against. A failed or skipped fetch and an unreachable Gitea are 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, noFetch bool) (repoCtx, error) {
	ctx := repoCtx{srcDir: srcDir, prs: map[string]agent.PullRequest{}}
	repo := filepath.Base(srcDir)
	if noFetch {
		ctx.unfetched = "fetch skipped"
		_, _ = fmt.Fprintf(out, "warn: fetch %s skipped (--no-fetch, remote state unverified)\n", repo)
	} else if err := agent.GitFetchPrune(srcDir, "origin", credentialHelperArgs()...); err != nil {
		ctx.unfetched = "fetch failed"
		_, _ = 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. Removal must never destroy state that
// exists nowhere else: a vanished working tree is the one case with nothing to
// lose, a locked or mid-rebase checkout holds sequencer state git itself refuses
// to discard, and a dirty checkout or a detached HEAD with unique commits holds
// work no branch is carrying. Past those guards, provably-upstream work loses its
// branch too, and anything unproven keeps its branch so no commit becomes
// 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}
	if wt.missing {
		res.verdict, res.reason = verdictRemove, "working tree gone, stale registration only"
		return res, nil
	}
	if wt.locked {
		res.verdict, res.reason = verdictKeep, "locked"
		return res, nil
	}
	op, err := agent.GitInProgressOp(wt.path)
	if err != nil {
		return res, err
	}
	if op != "" {
		res.verdict, res.reason = verdictKeep, op+" in progress"
		return res, nil
	}
	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
	}

	local, err := agent.GitCommitsNotOnRemotes(wt.path)
	if err != nil {
		return res, err
	}
	if wt.detached && local > 0 {
		res.verdict, res.reason = verdictKeep, fmt.Sprintf("detached HEAD carrying %s on no remote", commitCount(local))
		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, %s so %s is unverified", pr.Number, ctx.unfetched, 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, %s so %s is unverified", pr.Number, ctx.unfetched, 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"
	}
	if res.verdict == verdictRemove && local > 0 {
		res.reason += fmt.Sprintf(", %s on no remote so branch %s is kept", commitCount(local), wt.branch)
	}
	return res, nil
}

func commitCount(n int) string {
	if n == 1 {
		return "1 commit"
	}
	return fmt.Sprintf("%d commits", n)
}