fix(grep): grep context blowup from trajectory files

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,33 +185,46 @@ 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()
 
 	var matches []Match
 	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"
 }

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)
+	}
+}