gadfly/cmd/gadfly/tools.go

package main

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

	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
)

// 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
}

// 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(),
	}
}

// toolbox builds the reviewer's toolbox: the read-only repo tools, plus the
// delegate_investigation tool when a worker model is configured.
func (r *repoFS) toolbox() (*llm.Toolbox, error) {
	box := llm.NewToolbox("gadfly")
	tools := r.fsTools()
	if r.worker != nil {
		tools = append(tools, r.delegateTool())
	}
	for _, t := range tools {
		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
		},
	)
}

func (r *repoFS) getDiffTool() llm.Tool {
	return llm.DefineTool[struct{}](
		"get_diff",
		"Return the complete unified diff under review. The diff is also included (possibly truncated) in the task message; call this to get the full, untruncated text.",
		func(_ context.Context, _ struct{}) (any, error) {
			if strings.TrimSpace(r.diff) == "" {
				return "(empty diff)", nil
			}
			return r.diff, nil
		},
	)
}