From 3f1a91a8ed05053dfc8ff16931e9a3eea24b10ac Mon Sep 17 00:00:00 2001 From: huaiyj <8699003+huaiyj@user.noreply.gitee.com> Date: Mon, 24 Aug 2026 21:11:56 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E9=99=90=E5=88=B6=20ripgrep=20=E6=B5=81?= =?UTF-8?q?=E5=BC=8F=E8=BE=93=E5=87=BA=E5=86=85=E5=AD=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/ripgrep.go | 143 ++++++++++++++++++++++----------- internal/tools/ripgrep_test.go | 56 +++++++++++++ 2 files changed, 152 insertions(+), 47 deletions(-) create mode 100644 internal/tools/ripgrep_test.go diff --git a/internal/tools/ripgrep.go b/internal/tools/ripgrep.go index 6f7d6a0..bb10acb 100644 --- a/internal/tools/ripgrep.go +++ b/internal/tools/ripgrep.go @@ -1,16 +1,20 @@ package tools import ( + "bufio" "bytes" "context" "encoding/json" "fmt" + "io" "os/exec" "strings" ) const ( - DefaultMaxMatches = 20000 + DefaultMaxMatches = 20000 + maxRipgrepJSONLineSize = 4 * 1024 * 1024 + maxRipgrepStderrSize = 64 * 1024 ) // RipGrepMatch represents a single match from ripgrep @@ -159,71 +163,116 @@ func ripgrepHandler(ctx context.Context, params map[string]interface{}) (string, // Execute ripgrep cmd := exec.CommandContext(ctx, rgPath, args...) - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - - err = cmd.Run() - // Exit code 1 means no matches found, which is not an error + stdout, err := cmd.StdoutPipe() if err != nil { - if exitErr, ok := err.(*exec.ExitError); ok && exitErr.ExitCode() == 1 { - return fmt.Sprintf("No matches found for pattern \"%s\" in %s", p.Pattern, searchDir), nil - } - return "", fmt.Errorf("ripgrep failed: %w\nStderr: %s", err, stderr.String()) + return "", fmt.Errorf("failed to capture ripgrep stdout: %w", err) } - - // Parse ripgrep JSON output - matches, err := parseRipgrepJSON(stdout.Bytes()) + stderr, err := cmd.StderrPipe() if err != nil { - return "", fmt.Errorf("failed to parse ripgrep output: %w", err) + return "", fmt.Errorf("failed to capture ripgrep stderr: %w", err) + } + + var stderrOutput prefixBuffer + stderrOutput.limit = maxRipgrepStderrSize + if err := cmd.Start(); err != nil { + return "", fmt.Errorf("failed to start ripgrep: %w", err) } + stderrDone := make(chan struct{}) + go func() { + _, _ = io.Copy(&stderrOutput, stderr) + close(stderrDone) + }() + + matches, truncated, parseErr := readRipgrepMatches(stdout, DefaultMaxMatches) + if truncated || parseErr != nil { + _ = cmd.Process.Kill() + } + waitErr := cmd.Wait() + <-stderrDone - // Limit matches - if len(matches) > DefaultMaxMatches { - matches = matches[:DefaultMaxMatches] + if parseErr != nil { + return "", fmt.Errorf("failed to parse ripgrep output: %w", parseErr) + } + + // Exit code 1 means no matches found, which is not an error + if waitErr != nil && !truncated { + if exitErr, ok := waitErr.(*exec.ExitError); ok && exitErr.ExitCode() == 1 { + return fmt.Sprintf("No matches found for pattern \"%s\" in %s", p.Pattern, searchDir), nil + } + return "", fmt.Errorf("ripgrep failed: %w\nStderr: %s", waitErr, stderrOutput.String()) } // Format results - return formatRipgrepResults(matches, p.Pattern, searchDir, len(matches) >= DefaultMaxMatches), nil + return formatRipgrepResults(matches, p.Pattern, searchDir, truncated), nil } -// parseRipgrepJSON parses ripgrep JSON output -func parseRipgrepJSON(output []byte) ([]RipGrepMatch, error) { - var matches []RipGrepMatch - lines := bytes.Split(output, []byte("\n")) +func readRipgrepMatches(reader io.Reader, maxMatches int) ([]RipGrepMatch, bool, error) { + if maxMatches <= 0 { + return nil, true, nil + } - for _, line := range lines { - if len(line) == 0 { + matches := make([]RipGrepMatch, 0, min(maxMatches, 128)) + scanner := bufio.NewScanner(reader) + scanner.Buffer(make([]byte, 64*1024), maxRipgrepJSONLineSize) + for scanner.Scan() { + match, ok := parseRipgrepJSONLine(scanner.Bytes()) + if !ok { continue } - - var entry struct { - Type string `json:"type"` - Data struct { - Path struct { - Text string `json:"text"` - } `json:"path"` - Lines struct { - Text string `json:"text"` - } `json:"lines"` - LineNumber int `json:"line_number"` - } `json:"data"` + matches = append(matches, match) + if len(matches) >= maxMatches { + return matches, true, nil } + } + if err := scanner.Err(); err != nil { + return nil, false, err + } + return matches, false, nil +} - if err := json.Unmarshal(line, &entry); err != nil { - continue - } +func parseRipgrepJSONLine(line []byte) (RipGrepMatch, bool) { + var entry struct { + Type string `json:"type"` + Data struct { + Path struct { + Text string `json:"text"` + } `json:"path"` + Lines struct { + Text string `json:"text"` + } `json:"lines"` + LineNumber int `json:"line_number"` + } `json:"data"` + } - if entry.Type == "match" { - matches = append(matches, RipGrepMatch{ - FilePath: entry.Data.Path.Text, - LineNumber: entry.Data.LineNumber, - Line: strings.TrimRight(entry.Data.Lines.Text, "\n\r"), - }) + if err := json.Unmarshal(line, &entry); err != nil || entry.Type != "match" { + return RipGrepMatch{}, false + } + return RipGrepMatch{ + FilePath: entry.Data.Path.Text, + LineNumber: entry.Data.LineNumber, + Line: strings.TrimRight(entry.Data.Lines.Text, "\n\r"), + }, true +} + +type prefixBuffer struct { + buffer bytes.Buffer + limit int +} + +func (b *prefixBuffer) Write(p []byte) (int, error) { + originalLength := len(p) + remaining := b.limit - b.buffer.Len() + if remaining > 0 { + if len(p) > remaining { + p = p[:remaining] } + _, _ = b.buffer.Write(p) } + return originalLength, nil +} - return matches, nil +func (b *prefixBuffer) String() string { + return b.buffer.String() } // formatRipgrepResults formats the search results diff --git a/internal/tools/ripgrep_test.go b/internal/tools/ripgrep_test.go new file mode 100644 index 0000000..4f4a41a --- /dev/null +++ b/internal/tools/ripgrep_test.go @@ -0,0 +1,56 @@ +package tools + +import ( + "io" + "strings" + "testing" +) + +type countingReader struct { + reader io.Reader + read int +} + +func (r *countingReader) Read(p []byte) (int, error) { + n, err := r.reader.Read(p) + r.read += n + return n, err +} + +func TestReadRipgrepMatchesStopsAtLimit(t *testing.T) { + firstMatch := `{"type":"match","data":{"path":{"text":"a.go"},"lines":{"text":"hit\n"},"line_number":7}}` + input := firstMatch + "\n" + strings.Repeat("x", 1024*1024) + reader := &countingReader{reader: strings.NewReader(input)} + + matches, truncated, err := readRipgrepMatches(reader, 1) + + if err != nil { + t.Fatalf("readRipgrepMatches returned error: %v", err) + } + if !truncated { + t.Fatal("expected truncated result") + } + if len(matches) != 1 || matches[0].LineNumber != 7 { + t.Fatalf("matches = %#v, want one match at line 7", matches) + } + if reader.read >= len(input) { + t.Fatalf("reader consumed all %d bytes before truncating", reader.read) + } +} + +func TestPrefixBufferRetainsOnlyConfiguredPrefix(t *testing.T) { + var output prefixBuffer + output.limit = 5 + + n, err := output.Write([]byte("123456789")) + + if err != nil { + t.Fatalf("Write returned error: %v", err) + } + if n != 9 { + t.Fatalf("Write returned %d, want 9", n) + } + if got := output.String(); got != "12345" { + t.Fatalf("buffer = %q, want %q", got, "12345") + } +} -- Gitee