Vectrum
Sign In
238 lines · 6.6 KB · Go
package tools

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io/fs"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

// ripgrep is the path to rg, or empty when it is not installed. It is a
// variable so tests can exercise the pure-Go fallback on a machine that has rg.
var ripgrep = func() string {
	p, err := exec.LookPath("rg")
	if err != nil {
		return ""
	}
	return p
}()

// match is one hit, kept structured until output so results from either engine
// can be sorted the same way. ripgrep walks in parallel and returns files in a
// nondeterministic order; without this, the same search would produce different
// bytes on different machines, which would show up as eval noise and as prompt
// cache misses.
type match struct {
	path string
	line int
	text string
}

func (m match) String() string { return fmt.Sprintf("%s:%d: %s", m.path, m.line, m.text) }

func sortMatches(ms []match) {
	sort.Slice(ms, func(i, j int) bool {
		if ms[i].path != ms[j].path {
			return ms[i].path < ms[j].path
		}
		return ms[i].line < ms[j].line
	})
}

// GrepArgs is the grep input.
type GrepArgs struct {
	Pattern    string `json:"pattern" jsonschema_description:"Regular expression to search for"`
	Path       string `json:"path,omitempty" jsonschema_description:"File or directory to search; defaults to the working root"`
	Glob       string `json:"glob,omitempty" jsonschema_description:"Only search files matching this glob, e.g. *.go"`
	IgnoreCase bool   `json:"ignore_case,omitempty" jsonschema_description:"Match case-insensitively"`
}

// NewGrep builds the grep tool. It shells out to ripgrep when available and
// falls back to a pure-Go walk otherwise; the output format is the same either
// way, so the model cannot tell which ran. extra names additional directories
// that may be searched but never written, such as a skill's own directory.
func NewGrep(root string, extra ...string) Tool {
	return New("grep", strings.TrimSpace(`
Search file contents by regular expression. Returns "path:line: text" for each
match. Narrow with glob (e.g. "*.go") or path. This is how you find code you
have not been shown — reach for it before guessing at file names.`),
		func(ctx context.Context, in GrepArgs) (string, error) {
			if in.Pattern == "" {
				return "", fmt.Errorf("pattern is required")
			}
			re, err := compile(in.Pattern, in.IgnoreCase)
			if err != nil {
				return "", err
			}
			base, err := jailWith(root, extra, resolveBase(root, in.Path))
			if err != nil {
				return "", err
			}
			if _, err := os.Stat(base); err != nil {
				return "", fmt.Errorf("no such file or directory: %s", displayPath(root, base))
			}

			var matches []match
			if ripgrep != "" {
				matches, err = grepRipgrep(ctx, root, base, in)
			} else {
				matches, err = grepWalk(ctx, root, base, in, re)
			}
			if err != nil {
				return "", err
			}
			if len(matches) == 0 {
				where := displayPath(root, base)
				if in.Glob != "" {
					where += " matching " + in.Glob
				}
				return fmt.Sprintf("No matches for %q in %s.", in.Pattern, where), nil
			}

			sortMatches(matches)
			shown := matches
			if len(shown) > maxResults {
				shown = shown[:maxResults]
			}
			lines := make([]string, 0, len(shown))
			for _, m := range shown {
				lines = append(lines, m.String())
			}
			out := strings.Join(lines, "\n")
			if len(matches) > len(shown) {
				out += fmt.Sprintf("\n\n(%d of %d matches shown; narrow the pattern, path or glob to see the rest)",
					len(shown), len(matches))
			}
			return out, nil
		})
}

// compile validates the pattern up front so a bad regex comes back as a clear
// error the model can fix, rather than as ripgrep's exit status.
func compile(pattern string, ignoreCase bool) (*regexp.Regexp, error) {
	if ignoreCase {
		pattern = "(?i)" + pattern
	}
	re, err := regexp.Compile(pattern)
	if err != nil {
		return nil, fmt.Errorf("invalid regular expression: %w", err)
	}
	return re, nil
}

func grepRipgrep(ctx context.Context, root, base string, in GrepArgs) ([]match, error) {
	args := []string{"--line-number", "--no-heading", "--color=never", "--no-messages"}
	if in.IgnoreCase {
		args = append(args, "--ignore-case")
	}
	if in.Glob != "" {
		args = append(args, "--glob", in.Glob)
	}
	for dir := range skipDirs {
		args = append(args, "--glob", "!"+dir+"/")
	}
	args = append(args, "--regexp", in.Pattern, base)

	cmd := exec.CommandContext(ctx, ripgrep, args...)
	out, err := cmd.Output()
	if err != nil {
		// rg exits 1 when there are simply no matches; that is not a failure.
		var exitErr *exec.ExitError
		if !bytes.Contains(out, []byte(":")) && errorsAs(err, &exitErr) && exitErr.ExitCode() == 1 {
			return nil, nil
		}
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
		if !errorsAs(err, &exitErr) {
			return nil, fmt.Errorf("grep: %w", err)
		}
	}

	var matches []match
	sc := bufio.NewScanner(bytes.NewReader(out))
	sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
	for sc.Scan() {
		if m, ok := parseRipgrepLine(root, sc.Text()); ok {
			matches = append(matches, m)
		}
	}
	return matches, nil
}

// parseRipgrepLine splits "abs/path:12:text" into a match.
func parseRipgrepLine(root, line string) (match, bool) {
	parts := strings.SplitN(line, ":", 3)
	if len(parts) < 3 {
		return match{}, false
	}
	n, err := strconv.Atoi(parts[1])
	if err != nil {
		return match{}, false
	}
	return match{path: displayPath(root, parts[0]), line: n, text: capLine(strings.TrimSpace(parts[2]))}, true
}

func grepWalk(ctx context.Context, root, base string, in GrepArgs, re *regexp.Regexp) ([]match, error) {
	var matches []match
	err := filepath.WalkDir(base, func(p string, d fs.DirEntry, err error) error {
		if err != nil {
			return nil
		}
		if ctx.Err() != nil {
			return ctx.Err()
		}
		if d.IsDir() {
			if p != base && shouldSkip(d.Name(), in.Glob) {
				return fs.SkipDir
			}
			return nil
		}
		if in.Glob != "" {
			rel, relErr := filepath.Rel(base, p)
			if relErr != nil || !matchPath(in.Glob, rel) {
				return nil
			}
		}
		data, readErr := os.ReadFile(p)
		if readErr != nil || isBinary(data) {
			return nil
		}
		for i, line := range strings.Split(string(data), "\n") {
			if re.MatchString(line) {
				matches = append(matches, match{
					path: displayPath(root, p),
					line: i + 1,
					text: capLine(strings.TrimSpace(line)),
				})
			}
		}
		return nil
	})
	return matches, err
}

// isBinary treats a NUL byte in the first block as the signal, which is what
// grep and ripgrep both do.
func isBinary(data []byte) bool {
	if len(data) > 8000 {
		data = data[:8000]
	}
	return bytes.IndexByte(data, 0) >= 0
}

func errorsAs(err error, target **exec.ExitError) bool {
	e, ok := err.(*exec.ExitError)
	if ok {
		*target = e
	}
	return ok
}