From df3a69e3cd04cf70214c03c39fd770c77390ef7f Mon Sep 17 00:00:00 2001 From: GUOGUO <55723162+Dong1017@users.noreply.github.com> Date: Sat, 21 Mar 2026 17:40:04 +0800 Subject: [PATCH] Fix grep trajectory result blowup --- tools/fs/grep.go | 106 ++++++++++++++++++++++++++------ tools/fs/grep_test.go | 138 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 226 insertions(+), 18 deletions(-) create mode 100644 tools/fs/grep_test.go diff --git a/tools/fs/grep.go b/tools/fs/grep.go index bc27dd4..475823a 100644 --- a/tools/fs/grep.go +++ b/tools/fs/grep.go @@ -4,6 +4,7 @@ import ( "bufio" "context" "encoding/json" + "errors" "fmt" "os" "path/filepath" @@ -19,6 +20,16 @@ type GrepTool struct { workDir string } +const ( + defaultHeadLimit = 20 + maxHeadLimit = 100 + maxMatchLineRunes = 300 + matchLineTruncation = " [line truncated]" + resultSummaryTemplate = "%d matches" +) + +var errGrepLimitReached = errors.New("grep head limit reached") + // NewGrepTool creates a new grep tool. func NewGrepTool(workDir string) *GrepTool { return &GrepTool{workDir: workDir} @@ -55,6 +66,10 @@ func (t *GrepTool) Schema() llm.ToolSchema { Type: "boolean", Description: "Whether the search is case sensitive (default: true)", }, + "head_limit": { + Type: "integer", + Description: "Maximum number of matches to return (default: 20, max: 100)", + }, }, Required: []string{"pattern"}, } @@ -65,6 +80,7 @@ type grepParams struct { Path string `json:"path"` Include string `json:"include"` CaseSensitive bool `json:"case_sensitive"` + HeadLimit int `json:"head_limit"` } // Match represents a single grep match. @@ -82,11 +98,6 @@ func (t *GrepTool) Execute(ctx context.Context, params json.RawMessage) (*tools. return tools.ErrorResult(err), nil } - // Default case sensitive - if !p.CaseSensitive { - // Pattern will be handled with case-insensitive flag - } - // Resolve search path searchPath := "." if p.Path != "" { @@ -107,8 +118,10 @@ func (t *GrepTool) Execute(ctx context.Context, params json.RawMessage) (*tools. return tools.ErrorResultf("invalid pattern: %w", err), nil } + headLimit := normalizeHeadLimit(p.HeadLimit) + // Find files and search - matches, err := t.grep(ctx, fullPath, p.Include, re) + matches, truncated, err := t.grep(ctx, fullPath, p.Include, re, headLimit) if err != nil { return tools.ErrorResult(err), nil } @@ -121,30 +134,41 @@ func (t *GrepTool) Execute(ctx context.Context, params json.RawMessage) (*tools. var lines []string for _, m := range matches { relPath, _ := filepath.Rel(t.workDir, m.File) - lines = append(lines, fmt.Sprintf("%s:%d:%s", relPath, m.Line, m.Text)) + lines = append(lines, fmt.Sprintf("%s:%d:%s", relPath, m.Line, truncateMatchLine(m.Text))) } result := strings.Join(lines, "\n") - summary := fmt.Sprintf("%d matches", len(matches)) + summary := fmt.Sprintf(resultSummaryTemplate, len(matches)) + if truncated { + summary = fmt.Sprintf("%s (truncated at %d)", summary, headLimit) + result += fmt.Sprintf("\n[grep truncated after %d matches]", headLimit) + } return tools.StringResultWithSummary(result, summary), nil } -func (t *GrepTool) grep(ctx context.Context, root, include string, re *regexp.Regexp) ([]Match, error) { +func (t *GrepTool) grep(ctx context.Context, root, include string, re *regexp.Regexp, headLimit int) ([]Match, bool, error) { var matches []Match info, err := os.Stat(root) if err != nil { - return nil, err + return nil, false, err } if !info.IsDir() { // Single file - return t.searchFile(root, re) + fileMatches, limitReached, err := t.searchFile(root, re, headLimit) + if err != nil { + return nil, false, err + } + return fileMatches, limitReached, nil } // Walk directory err = filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { + if ctx.Err() != nil { + return ctx.Err() + } if err != nil { return nil // Skip errors } @@ -161,26 +185,38 @@ func (t *GrepTool) grep(ctx context.Context, root, include string, re *regexp.Re } } - fileMatches, err := t.searchFile(path, re) + if shouldSkipTrajectory(path, d) { + return nil + } + + remaining := headLimit - len(matches) + if remaining <= 0 { + return errGrepLimitReached + } + + fileMatches, limitReached, err := t.searchFile(path, re, remaining) if err != nil { return nil // Skip file errors } matches = append(matches, fileMatches...) + if limitReached || len(matches) >= headLimit { + return errGrepLimitReached + } return nil }) - if err != nil { - return nil, err + if err != nil && !errors.Is(err, errGrepLimitReached) { + return nil, false, err } - return matches, nil + return matches, errors.Is(err, errGrepLimitReached), nil } -func (t *GrepTool) searchFile(path string, re *regexp.Regexp) ([]Match, error) { +func (t *GrepTool) searchFile(path string, re *regexp.Regexp, headLimit int) ([]Match, bool, error) { file, err := os.Open(path) if err != nil { - return nil, err + return nil, false, err } defer file.Close() @@ -188,6 +224,7 @@ func (t *GrepTool) searchFile(path string, re *regexp.Regexp) ([]Match, error) { scanner := bufio.NewScanner(file) scanner.Buffer(make([]byte, 64*1024), 1024*1024) lineNum := 0 + limitReached := false for scanner.Scan() { lineNum++ @@ -200,8 +237,41 @@ func (t *GrepTool) searchFile(path string, re *regexp.Regexp) ([]Match, error) { Column: loc[0] + 1, Text: line, }) + if len(matches) >= headLimit { + limitReached = true + break + } } } - return matches, scanner.Err() + return matches, limitReached, scanner.Err() +} + +func normalizeHeadLimit(limit int) int { + switch { + case limit <= 0: + return defaultHeadLimit + case limit > maxHeadLimit: + return maxHeadLimit + default: + return limit + } +} + +func truncateMatchLine(line string) string { + runes := []rune(line) + if len(runes) <= maxMatchLineRunes { + return line + } + return string(runes[:maxMatchLineRunes]) + matchLineTruncation +} + +func shouldSkipTrajectory(path string, d os.DirEntry) bool { + if d.IsDir() { + return false + } + if !strings.HasSuffix(d.Name(), ".trajectory.jsonl") { + return false + } + return filepath.Base(filepath.Dir(path)) == ".cache" } diff --git a/tools/fs/grep_test.go b/tools/fs/grep_test.go new file mode 100644 index 0000000..0f9c69c --- /dev/null +++ b/tools/fs/grep_test.go @@ -0,0 +1,138 @@ +package fs + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestGrepSkipsTrajectoryDuringDirectoryWalk(t *testing.T) { + workDir := t.TempDir() + writeTestFile(t, workDir, "notes.txt", "skill in project\n") + writeTestFile(t, workDir, ".cache/session.trajectory.jsonl", `{"message":"skill in trace"}`+"\n") + + tool := NewGrepTool(workDir) + result := executeGrep(t, tool, grepParams{Pattern: "skill", Path: "."}) + + if strings.Contains(result.Content, ".cache/session.trajectory.jsonl") { + t.Fatalf("directory walk should skip trajectory files, got %q", result.Content) + } + if !strings.Contains(result.Content, "notes.txt:1:skill in project") { + t.Fatalf("expected normal file match, got %q", result.Content) + } +} + +func TestGrepAllowsExplicitTrajectoryFilePath(t *testing.T) { + workDir := t.TempDir() + writeTestFile(t, workDir, ".cache/session.trajectory.jsonl", `{"message":"skill in trace"}`+"\n") + + tool := NewGrepTool(workDir) + result := executeGrep(t, tool, grepParams{ + Pattern: "skill", + Path: ".cache/session.trajectory.jsonl", + }) + + if !strings.Contains(filepath.ToSlash(result.Content), `.cache/session.trajectory.jsonl:1:{"message":"skill in trace"}`) { + t.Fatalf("explicit trajectory file path should be searchable, got %q", result.Content) + } +} + +func TestGrepHeadLimitTruncatesMatches(t *testing.T) { + workDir := t.TempDir() + writeTestFile(t, workDir, "many.txt", "match 1\nmatch 2\nmatch 3\n") + + tool := NewGrepTool(workDir) + result := executeGrep(t, tool, grepParams{ + Pattern: "match", + Path: ".", + HeadLimit: 2, + }) + + if strings.Contains(result.Content, "many.txt:3:match 3") { + t.Fatalf("expected head limit to stop after 2 matches, got %q", result.Content) + } + if !strings.Contains(result.Content, "[grep truncated after 2 matches]") { + t.Fatalf("expected truncation marker, got %q", result.Content) + } + if result.Summary != "2 matches (truncated at 2)" { + t.Fatalf("summary = %q, want %q", result.Summary, "2 matches (truncated at 2)") + } +} + +func TestGrepTruncatesLongMatchedLine(t *testing.T) { + workDir := t.TempDir() + longTail := strings.Repeat("x", maxMatchLineRunes+50) + writeTestFile(t, workDir, "long.txt", "match "+longTail+"\n") + + tool := NewGrepTool(workDir) + result := executeGrep(t, tool, grepParams{Pattern: "match", Path: "."}) + + if !strings.Contains(result.Content, matchLineTruncation) { + t.Fatalf("expected long line truncation marker, got %q", result.Content) + } + if strings.Contains(result.Content, longTail) { + t.Fatalf("expected long line content to be truncated, got %q", result.Content) + } +} + +func TestGrepNoMatchesFound(t *testing.T) { + workDir := t.TempDir() + writeTestFile(t, workDir, "notes.txt", "no relevant content\n") + + tool := NewGrepTool(workDir) + result := executeGrep(t, tool, grepParams{Pattern: "skill", Path: "."}) + + if result.Content != "No matches found" { + t.Fatalf("content = %q, want %q", result.Content, "No matches found") + } + if result.Summary != "0 matches" { + t.Fatalf("summary = %q, want %q", result.Summary, "0 matches") + } +} + +func executeGrep(t *testing.T, tool *GrepTool, params grepParams) *struct { + Content string + Summary string +} { + t.Helper() + + raw, err := json.Marshal(params) + if err != nil { + t.Fatalf("Marshal params failed: %v", err) + } + + result, err := tool.Execute(context.Background(), raw) + if err != nil { + t.Fatalf("Execute failed: %v", err) + } + if result.Error != nil { + t.Fatalf("tool returned error: %v", result.Error) + } + + return &struct { + Content string + Summary string + }{ + Content: result.Content, + Summary: result.Summary, + } +} + +func writeTestFile(t *testing.T, workDir, relPath, content string) { + t.Helper() + + fullPath, err := resolveSafePath(workDir, relPath) + if err != nil { + t.Fatalf("resolveSafePath failed: %v", err) + } + dir := filepath.Dir(fullPath) + if err := os.MkdirAll(dir, 0755); err != nil { + t.Fatalf("MkdirAll failed: %v", err) + } + if err := os.WriteFile(fullPath, []byte(content), 0644); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } +}