fix(grep): do not go outside cwd, add timeout (#2188)

Carlos Alexandro Becker created 1 day ago

* fix(grep): do not go outside cwd, add timeout

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

* fix: timeout config

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

---------

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

Change summary

internal/agent/agentic_fetch_tool.go |  2 +-
internal/agent/common_test.go        |  2 +-
internal/agent/coordinator.go        |  2 +-
internal/agent/tools/grep.go         |  9 ++++++---
internal/config/config.go            | 13 ++++++++++++-
5 files changed, 21 insertions(+), 7 deletions(-)

Detailed changes

internal/agent/agentic_fetch_tool.go

@@ -167,7 +167,7 @@ func (c *coordinator) agenticFetchTool(_ context.Context, client *http.Client) (
 				webFetchTool,
 				webSearchTool,
 				tools.NewGlobTool(tmpDir),
-				tools.NewGrepTool(tmpDir),
+				tools.NewGrepTool(tmpDir, c.cfg.Tools.Grep),
 				tools.NewSourcegraphTool(client),
 				tools.NewViewTool(c.lspManager, c.permissions, c.filetracker, tmpDir),
 			}

internal/agent/common_test.go

@@ -208,7 +208,7 @@ func coderAgent(r *vcr.Recorder, env fakeEnv, large, small fantasy.LanguageModel
 		tools.NewMultiEditTool(nil, env.permissions, env.history, *env.filetracker, env.workingDir),
 		tools.NewFetchTool(env.permissions, env.workingDir, r.GetDefaultClient()),
 		tools.NewGlobTool(env.workingDir),
-		tools.NewGrepTool(env.workingDir),
+		tools.NewGrepTool(env.workingDir, cfg.Tools.Grep),
 		tools.NewLsTool(env.permissions, env.workingDir, cfg.Tools.Ls),
 		tools.NewSourcegraphTool(r.GetDefaultClient()),
 		tools.NewViewTool(nil, env.permissions, *env.filetracker, env.workingDir),

internal/agent/coordinator.go

@@ -426,7 +426,7 @@ func (c *coordinator) buildTools(ctx context.Context, agent config.Agent) ([]fan
 		tools.NewMultiEditTool(c.lspManager, c.permissions, c.history, c.filetracker, c.cfg.WorkingDir()),
 		tools.NewFetchTool(c.permissions, c.cfg.WorkingDir(), nil),
 		tools.NewGlobTool(c.cfg.WorkingDir()),
-		tools.NewGrepTool(c.cfg.WorkingDir()),
+		tools.NewGrepTool(c.cfg.WorkingDir(), c.cfg.Tools.Grep),
 		tools.NewLsTool(c.permissions, c.cfg.WorkingDir(), c.cfg.Tools.Ls),
 		tools.NewSourcegraphTool(nil),
 		tools.NewTodosTool(c.sessions),

internal/agent/tools/grep.go

@@ -18,6 +18,7 @@ import (
 	"time"
 
 	"charm.land/fantasy"
+	"github.com/charmbracelet/crush/internal/config"
 	"github.com/charmbracelet/crush/internal/csync"
 	"github.com/charmbracelet/crush/internal/fsext"
 )
@@ -100,7 +101,7 @@ func escapeRegexPattern(pattern string) string {
 	return escaped
 }
 
-func NewGrepTool(workingDir string) fantasy.AgentTool {
+func NewGrepTool(workingDir string, config config.ToolGrep) fantasy.AgentTool {
 	return fantasy.NewAgentTool(
 		GrepToolName,
 		string(grepDescription),
@@ -109,7 +110,6 @@ func NewGrepTool(workingDir string) fantasy.AgentTool {
 				return fantasy.NewTextErrorResponse("pattern is required"), nil
 			}
 
-			// If literal_text is true, escape the pattern
 			searchPattern := params.Pattern
 			if params.LiteralText {
 				searchPattern = escapeRegexPattern(params.Pattern)
@@ -120,7 +120,10 @@ func NewGrepTool(workingDir string) fantasy.AgentTool {
 				searchPath = workingDir
 			}
 
-			matches, truncated, err := searchFiles(ctx, searchPattern, searchPath, params.Include, 100)
+			searchCtx, cancel := context.WithTimeout(ctx, config.GetTimeout())
+			defer cancel()
+
+			matches, truncated, err := searchFiles(searchCtx, searchPattern, searchPath, params.Include, 100)
 			if err != nil {
 				return fantasy.NewTextErrorResponse(fmt.Sprintf("error searching files: %v", err)), nil
 			}

internal/config/config.go

@@ -350,7 +350,8 @@ type Agent struct {
 }
 
 type Tools struct {
-	Ls ToolLs `json:"ls,omitzero"`
+	Ls   ToolLs   `json:"ls,omitzero"`
+	Grep ToolGrep `json:"grep,omitzero"`
 }
 
 type ToolLs struct {
@@ -358,10 +359,20 @@ type ToolLs struct {
 	MaxItems *int `json:"max_items,omitempty" jsonschema:"description=Maximum number of items to return for the ls tool,default=1000,example=100"`
 }
 
+// Limits returns the user-defined max-depth and max-items, or their defaults.
 func (t ToolLs) Limits() (depth, items int) {
 	return ptrValOr(t.MaxDepth, 0), ptrValOr(t.MaxItems, 0)
 }
 
+type ToolGrep struct {
+	Timeout *time.Duration `json:"timeout,omitempty" jsonschema:"description=Timeout for the grep tool call,default=5s,example=10s"`
+}
+
+// GetTimeout returns the user-defined timeout or the default.
+func (t ToolGrep) GetTimeout() time.Duration {
+	return ptrValOr(t.Timeout, 5*time.Second)
+}
+
 // Config holds the configuration for crush.
 type Config struct {
 	Schema string `json:"$schema,omitempty"`