gadfly/cmd/gadfly/tools.go

package main

import (
	"bufio"
	"context"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"

	llm "gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)

// Tool output bounds. The reviewer is a chat agent with a finite context, so
// every tool caps how much it can pull in one call — a runaway read_file or
// grep would blow the window and stall the loop.
const (
	maxFileBytes    = 64 * 1024 // per read_file call
	maxReadLines    = 800       // per read_file call
	maxGrepResults  = 200       // per grep call
	maxFindResults  = 200       // per find_files call
	maxLineLen      = 400       // truncate any single returned line to this
	maxGetDiffLines = 800       // per get_diff call (paginated window)
	maxGetDiffBytes = 64 * 1024 // per get_diff call
)

// skipDirs are never descended into by grep / find_files — noise and bulk that
// a code reviewer never needs and that would swamp the results.
var skipDirs = map[string]bool{
	".git":         true,
	"node_modules": true,
	"vendor":       true,
}

// repoFS is a read-only, sandboxed view of the checked-out repository. Every
// path argument from the model is resolved against root and rejected if it
// escapes (symlink or `..` traversal), so a hostile diff can never make the
// reviewer read outside the checkout.
type repoFS struct {
	root   string    // absolute, symlink-resolved repo root
	diff   string    // the full PR unified diff (served by get_diff)
	worker llm.Model // optional cheap model for delegate_investigation; nil = no delegation

	// diffLines caches the split diff so paging through get_diff doesn't re-split
	// the whole (possibly large) diff on every call. The diff is immutable, so the
	// cache is computed once and safe to share across concurrent lenses.
	diffOnce  sync.Once
	diffLines []string
}

// newRepoFS resolves root to an absolute, symlink-free path.
func newRepoFS(root, diff string) (*repoFS, error) {
	abs, err := filepath.Abs(root)
	if err != nil {
		return nil, fmt.Errorf("resolve repo dir: %w", err)
	}
	// EvalSymlinks so prefix containment checks survive a symlinked root
	// (e.g. macOS /tmp -> /private/tmp).
	if resolved, err := filepath.EvalSymlinks(abs); err == nil {
		abs = resolved
	}
	info, err := os.Stat(abs)
	if err != nil {
		return nil, fmt.Errorf("repo dir %q: %w", root, err)
	}
	if !info.IsDir() {
		return nil, fmt.Errorf("repo dir %q is not a directory", root)
	}
	return &repoFS{root: abs, diff: diff}, nil
}

// resolve maps a model-supplied relative path to an absolute path inside the
// sandbox, rejecting anything that escapes root. An empty path means root.
func (r *repoFS) resolve(rel string) (string, error) {
	rel = strings.TrimSpace(rel)
	rel = strings.TrimPrefix(rel, "./")
	if rel == "" || rel == "." {
		return r.root, nil
	}
	if filepath.IsAbs(rel) {
		// Allow an absolute path only if it already points inside the sandbox.
		clean := filepath.Clean(rel)
		if err := r.contains(clean); err != nil {
			return "", err
		}
		return clean, nil
	}
	joined := filepath.Clean(filepath.Join(r.root, rel))
	if err := r.contains(joined); err != nil {
		return "", err
	}
	return joined, nil
}

// contains verifies abs is root or lives beneath it.
func (r *repoFS) contains(abs string) error {
	if abs == r.root {
		return nil
	}
	if !strings.HasPrefix(abs, r.root+string(os.PathSeparator)) {
		return fmt.Errorf("path escapes the repository sandbox")
	}
	return nil
}

// fsTools is the set of read-only repository tools.
func (r *repoFS) fsTools() []llm.Tool {
	return []llm.Tool{
		r.readFileTool(),
		r.listDirTool(),
		r.grepTool(),
		r.findFilesTool(),
		r.getDiffTool(),
	}
}

// allTools is the single source of truth for the reviewer's tool set: the
// read-only repo tools, plus delegate_investigation when a worker model is
// configured. Both the executus registry (gadflyToolRegistry, the production
// path) and toolbox() build from this list.
func (r *repoFS) allTools() []llm.Tool {
	tools := r.fsTools()
	if r.worker != nil {
		tools = append(tools, r.delegateTool())
	}
	return tools
}

// toolbox builds a majordomo toolbox from allTools(). The production review path
// now goes through executus's tool.Registry (see executus.go); this remains for
// the toolbox-level tests (the `call` helper).
func (r *repoFS) toolbox() (*llm.Toolbox, error) {
	box := llm.NewToolbox("gadfly")
	for _, t := range r.allTools() {
		if err := box.Add(t); err != nil {
			return nil, fmt.Errorf("add tool %q: %w", t.Name, err)
		}
	}
	return box, nil
}

// workerToolbox is the toolbox handed to a delegated worker sub-agent: the
// read-only repo tools only (no delegate tool — workers don't sub-delegate).
func (r *repoFS) workerToolbox() (*llm.Toolbox, error) {
	box := llm.NewToolbox("gadfly-worker")
	for _, t := range r.fsTools() {
		if err := box.Add(t); err != nil {
			return nil, fmt.Errorf("add tool %q: %w", t.Name, err)
		}
	}
	return box, nil
}

type readFileArgs struct {
	Path      string `json:"path" description:"Repository-relative path of the file to read, e.g. pkg/logic/agentexec/pipeline.go"`
	StartLine int    `json:"start_line,omitempty" description:"Optional 1-based line to start from (default 1)."`
	Limit     int    `json:"limit,omitempty" description:"Optional max number of lines to return (default/maximum 800)."`
}

func (r *repoFS) readFileTool() llm.Tool {
	return llm.DefineTool[readFileArgs](
		"read_file",
		"Read a file from the repository at its current checked-out state, with line numbers. Use this to verify the surrounding code, imports, and symbols a diff hunk touches before reporting an issue.",
		func(_ context.Context, args readFileArgs) (any, error) {
			abs, err := r.resolve(args.Path)
			if err != nil {
				return nil, err
			}
			info, err := os.Stat(abs)
			if err != nil {
				return nil, fmt.Errorf("stat %q: %w", args.Path, err)
			}
			if info.IsDir() {
				return nil, fmt.Errorf("%q is a directory; use list_dir", args.Path)
			}
			f, err := os.Open(abs)
			if err != nil {
				return nil, fmt.Errorf("open %q: %w", args.Path, err)
			}
			defer f.Close()

			start := args.StartLine
			if start < 1 {
				start = 1
			}
			limit := args.Limit
			if limit <= 0 || limit > maxReadLines {
				limit = maxReadLines
			}

			var b strings.Builder
			sc := bufio.NewScanner(f)
			sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
			lineNo := 0
			emitted := 0
			for sc.Scan() {
				lineNo++
				if lineNo < start {
					continue
				}
				if emitted >= limit || b.Len() >= maxFileBytes {
					fmt.Fprintf(&b, "... (truncated at line %d; call read_file again with start_line=%d for more)\n", lineNo, lineNo)
					break
				}
				line := sc.Text()
				if len(line) > maxLineLen {
					line = line[:maxLineLen] + "…"
				}
				fmt.Fprintf(&b, "%d\t%s\n", lineNo, line)
				emitted++
			}
			if err := sc.Err(); err != nil {
				return nil, fmt.Errorf("read %q: %w", args.Path, err)
			}
			if emitted == 0 {
				return fmt.Sprintf("(%s has no lines at/after %d; file has %d lines)", args.Path, start, lineNo), nil
			}
			return b.String(), nil
		},
	)
}

type listDirArgs struct {
	Path string `json:"path,omitempty" description:"Optional repository-relative directory (default: repo root)."`
}

func (r *repoFS) listDirTool() llm.Tool {
	return llm.DefineTool[listDirArgs](
		"list_dir",
		"List the entries of a directory in the repository (directories marked with a trailing /). Use it to discover where code lives before reading.",
		func(_ context.Context, args listDirArgs) (any, error) {
			abs, err := r.resolve(args.Path)
			if err != nil {
				return nil, err
			}
			entries, err := os.ReadDir(abs)
			if err != nil {
				return nil, fmt.Errorf("list %q: %w", args.Path, err)
			}
			names := make([]string, 0, len(entries))
			for _, e := range entries {
				name := e.Name()
				if e.IsDir() {
					name += "/"
				}
				names = append(names, name)
			}
			sort.Strings(names)
			if len(names) == 0 {
				return "(empty directory)", nil
			}
			return strings.Join(names, "\n"), nil
		},
	)
}

type grepArgs struct {
	Pattern    string `json:"pattern" description:"A Go (RE2) regular expression to search for."`
	Path       string `json:"path,omitempty" description:"Optional repository-relative file or subdirectory to scope the search (default: whole repo)."`
	MaxResults int    `json:"max_results,omitempty" description:"Optional cap on matching lines returned (default/maximum 200)."`
}

func (r *repoFS) grepTool() llm.Tool {
	return llm.DefineTool[grepArgs](
		"grep",
		"Search the repository's text files for a regular expression and return matching `path:line: text`. Use it to check whether a symbol, import, or call exists elsewhere before claiming a cross-file problem.",
		func(_ context.Context, args grepArgs) (any, error) {
			if strings.TrimSpace(args.Pattern) == "" {
				return nil, fmt.Errorf("pattern is required")
			}
			re, err := regexp.Compile(args.Pattern)
			if err != nil {
				return nil, fmt.Errorf("invalid regexp: %w", err)
			}
			base, err := r.resolve(args.Path)
			if err != nil {
				return nil, err
			}
			limit := args.MaxResults
			if limit <= 0 || limit > maxGrepResults {
				limit = maxGrepResults
			}

			var out []string
			truncated := false
			walkErr := filepath.WalkDir(base, func(path string, d os.DirEntry, err error) error {
				if err != nil {
					return nil // skip unreadable entries
				}
				if d.IsDir() {
					if skipDirs[d.Name()] && path != base {
						return filepath.SkipDir
					}
					return nil
				}
				if len(out) >= limit {
					truncated = true
					return filepath.SkipAll
				}
				matchesInFile(path, r.root, re, limit, &out)
				return nil
			})
			if walkErr != nil {
				return nil, fmt.Errorf("search: %w", walkErr)
			}
			if len(out) > limit {
				out = out[:limit]
				truncated = true
			}
			if len(out) == 0 {
				return "(no matches)", nil
			}
			res := strings.Join(out, "\n")
			if truncated {
				res += fmt.Sprintf("\n... (truncated at %d matches; narrow the pattern or path)", limit)
			}
			return res, nil
		},
	)
}

// matchesInFile appends "relpath:line: text" for each regexp match in a single
// text file, stopping once the global cap is reached. Binary files (NUL in the
// first chunk) and oversized files are skipped.
func matchesInFile(path, root string, re *regexp.Regexp, limit int, out *[]string) {
	f, err := os.Open(path)
	if err != nil {
		return
	}
	defer f.Close()

	rel, relErr := filepath.Rel(root, path)
	if relErr != nil {
		rel = path
	}
	sc := bufio.NewScanner(f)
	sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
	lineNo := 0
	for sc.Scan() {
		if len(*out) >= limit {
			return
		}
		lineNo++
		line := sc.Text()
		if lineNo == 1 && strings.IndexByte(line, 0) >= 0 {
			return // looks binary
		}
		if re.MatchString(line) {
			trimmed := strings.TrimSpace(line)
			if len(trimmed) > maxLineLen {
				trimmed = trimmed[:maxLineLen] + "…"
			}
			*out = append(*out, fmt.Sprintf("%s:%d: %s", rel, lineNo, trimmed))
		}
	}
}

type findFilesArgs struct {
	Name       string `json:"name" description:"Case-insensitive substring of the file path to match, e.g. \"pipeline.go\" or \"agentexec/\"."`
	MaxResults int    `json:"max_results,omitempty" description:"Optional cap on paths returned (default/maximum 200)."`
}

func (r *repoFS) findFilesTool() llm.Tool {
	return llm.DefineTool[findFilesArgs](
		"find_files",
		"Find files whose repository-relative path contains a case-insensitive substring. Use it to locate a file by name when you don't know its directory.",
		func(_ context.Context, args findFilesArgs) (any, error) {
			needle := strings.ToLower(strings.TrimSpace(args.Name))
			if needle == "" {
				return nil, fmt.Errorf("name is required")
			}
			limit := args.MaxResults
			if limit <= 0 || limit > maxFindResults {
				limit = maxFindResults
			}
			var out []string
			truncated := false
			_ = filepath.WalkDir(r.root, func(path string, d os.DirEntry, err error) error {
				if err != nil {
					return nil
				}
				if d.IsDir() {
					if skipDirs[d.Name()] && path != r.root {
						return filepath.SkipDir
					}
					return nil
				}
				if len(out) >= limit {
					truncated = true
					return filepath.SkipAll
				}
				rel, relErr := filepath.Rel(r.root, path)
				if relErr != nil {
					return nil
				}
				if strings.Contains(strings.ToLower(rel), needle) {
					out = append(out, rel)
				}
				return nil
			})
			sort.Strings(out)
			if len(out) == 0 {
				return "(no files matched)", nil
			}
			res := strings.Join(out, "\n")
			if truncated {
				res += fmt.Sprintf("\n... (truncated at %d files; narrow the name)", limit)
			}
			return res, nil
		},
	)
}

type getDiffArgs struct {
	Path      string `json:"path,omitempty" description:"Optional changed-file path (e.g. pkg/foo/bar.go); returns ONLY that file's diff hunks. Omit for the whole diff. Use this on a large PR to pull just the file a finding is about."`
	StartLine int    `json:"start_line,omitempty" description:"Optional 1-based line to start from within the (whole or path-scoped) diff (default 1)."`
	Limit     int    `json:"limit,omitempty" description:"Optional max number of diff lines to return (default/maximum 800)."`
}

func (r *repoFS) getDiffTool() llm.Tool {
	return llm.DefineTool[getDiffArgs](
		"get_diff",
		"Return the unified diff under review as a numbered, PAGINATED window (like read_file) — not the whole diff at once, so a huge PR can't blow the context window. Pass `path` to fetch just one changed file's hunks, or `start_line`/`limit` to page through. A truncated copy of the diff is also embedded in the task message.",
		func(_ context.Context, args getDiffArgs) (any, error) {
			scope := "the diff"
			var lines []string
			if p := strings.TrimSpace(args.Path); p != "" {
				lines = diffLinesForPath(r.diff, p)
				if len(lines) == 0 {
					return fmt.Sprintf("(no diff hunks for path %q; check it against the changed-files list — the path must match a file the diff touches)", p), nil
				}
				scope = "the diff for " + p
			} else {
				lines = r.diffAllLines()
				if len(lines) == 0 {
					return "(empty diff)", nil
				}
			}
			return windowDiff(lines, scope, args.StartLine, args.Limit), nil
		},
	)
}

// diffAllLines returns the whole diff split into lines, cached so paging never
// re-splits the (possibly large) diff.
func (r *repoFS) diffAllLines() []string {
	r.diffOnce.Do(func() { r.diffLines = splitDiffLines(r.diff) })
	return r.diffLines
}

// splitDiffLines splits a unified diff into lines, dropping the single trailing
// empty element a trailing newline produces — otherwise windowDiff would emit a
// blank final line and over-count the total by one.
func splitDiffLines(diff string) []string {
	lines := strings.Split(diff, "\n")
	if n := len(lines); n > 0 && lines[n-1] == "" {
		lines = lines[:n-1]
	}
	return lines
}

// windowDiff returns a numbered, paginated slice of pre-split diff lines,
// mirroring read_file's caps (maxGetDiffLines / maxGetDiffBytes / maxLineLen) so
// a single get_diff call can never dump a multi-hundred-KB diff into the
// transcript — the amplifier behind the large-PR token burn. The full diff stays
// reachable by paging with start_line, or scoped per file via the path arg.
func windowDiff(lines []string, scope string, start, limit int) string {
	total := len(lines)
	if start < 1 {
		start = 1
	}
	if limit <= 0 || limit > maxGetDiffLines {
		limit = maxGetDiffLines
	}
	if start > total {
		return fmt.Sprintf("(%s has %d lines; nothing at/after line %d)", scope, total, start)
	}
	var b strings.Builder
	emitted := 0
	i := start - 1
	for ; i < total; i++ {
		if emitted >= limit || b.Len() >= maxGetDiffBytes {
			break
		}
		line := lines[i]
		if len(line) > maxLineLen {
			line = line[:maxLineLen] + "…"
		}
		fmt.Fprintf(&b, "%d\t%s\n", i+1, line)
		emitted++
	}
	if i < total {
		// i (0-based) is the first line NOT emitted; line i was the last shown.
		fmt.Fprintf(&b, "... (%s truncated after line %d of %d; call get_diff again with start_line=%d for the rest, or pass a `path` to scope to one file)\n", scope, i, total, i+1)
	}
	return b.String()
}

// diffLinesForPath returns the unified-diff lines for one changed file: from the
// `diff --git` header that names path through to the next header (or end). The
// header names two path tokens (a/<old> b/<new>); a match is on a WHOLE token,
// so path "foo.go" does not pull in "barfoo.go" — and a trailing "/" scopes to a
// directory (e.g. "pkg/foo/" matches pkg/foo/bar.go).
func diffLinesForPath(diff, path string) []string {
	want := strings.TrimPrefix(strings.TrimPrefix(strings.TrimSpace(path), "a/"), "b/")
	var out []string
	inSection := false
	for _, ln := range splitDiffLines(diff) {
		if strings.HasPrefix(ln, "diff --git ") {
			inSection = diffHeaderNames(ln, want)
		}
		if inSection {
			out = append(out, ln)
		}
	}
	return out
}

// diffHeaderNames reports whether a `diff --git a/X b/Y` header names want as one
// of its (a/b-stripped) path tokens — exact whole-token match, or a directory
// prefix when want ends with "/".
func diffHeaderNames(header, want string) bool {
	fields := strings.Fields(header)
	if len(fields) < 3 {
		return false
	}
	for _, f := range fields[2:] {
		p := strings.TrimPrefix(strings.TrimPrefix(f, "a/"), "b/")
		if p == want || (strings.HasSuffix(want, "/") && strings.HasPrefix(p, want)) {
			return true
		}
	}
	return false
}