1package tools
2
3import (
4 "bufio"
5 "bytes"
6 "context"
7 "encoding/json"
8 "errors"
9 "fmt"
10 "go/format"
11 "os"
12 "os/exec"
13 "path/filepath"
14 "regexp"
15 "sort"
16 "strings"
17 "time"
18 "unicode/utf8"
19
20 "github.com/bnema/zut/packages/agent/lsp"
21 "github.com/bnema/zut/packages/core"
22 "github.com/bnema/zut/packages/provider"
23)
24
25const (
26 maxASTOutputBytes = 60 * 1024
27 maxASTCaptureBytes = 4 * 1024 * 1024
28)
29
30// ASTTool performs structural searches and rewrites through ast-grep.
31type ASTTool struct {
32 CWD string
33 Sandbox *Sandbox
34 LSP *lsp.Manager
35 LSPDiagnostics bool
36 ReadOnly bool
37 LookPath func(string) (string, error)
38}
39
40var (
41 _ core.Tool = (*ASTTool)(nil)
42 _ core.ToolPreviewer = (*ASTTool)(nil)
43)
44
45type astArgs struct {
46 Pattern string `json:"pattern"`
47 Language string `json:"language"`
48 Path string `json:"path"`
49 Rewrite string `json:"rewrite,omitempty"`
50}
51
52const astSchema = `{"type":"object","properties":{"pattern":{"type":"string","description":"Code-shaped ast-grep pattern. $NAME matches one node, $$$NAME matches zero or more nodes (arguments, statements)."},"language":{"type":"string","description":"Optional ast-grep language id (go, rust, python, typescript, tsx, javascript, ...). Omit to infer it from each file's extension."},"path":{"type":"string","description":"File or directory to search, relative to the workspace or absolute."},"rewrite":{"type":"string","description":"Optional replacement reusing captured metavariables. Zut previews the diff and applies every match."}},"required":["pattern","path"],"additionalProperties":false}`
53
54func (t *ASTTool) Name() string { return "ast" }
55func (t *ASTTool) Description() string {
56 return "Structural code search and rewrite (ast-grep). Prefer it over grep/edit when looking for code shapes: calls, definitions, statements, or a rename/refactor across many places. " +
57 "Matches ignore formatting and span multiple lines. Examples: " +
58 "calls `fmt.Errorf($MSG, $$$REST)`; " +
59 "Go methods `func ($R $T) $NAME($$$) $$$ { $$$ }`; " +
60 "error checks `if err != nil { $$$ }`; " +
61 "JS `console.log($$$)`; " +
62 "rename: pattern `oldFn($$$A)` with rewrite `newFn($$$A)`. " +
63 "Use grep for plain text, comments, or strings."
64}
65func (t *ASTTool) Schema() json.RawMessage { return json.RawMessage(astSchema) }
66
67func (t *ASTTool) Preview(ctx context.Context, raw json.RawMessage) (core.ToolResult, error) {
68 args, err := parseASTArgs(raw)
69 if err != nil {
70 return core.ToolResult{}, err
71 }
72 plan, err := t.plan(ctx, args)
73 if err != nil {
74 return core.ToolResult{}, err
75 }
76 return plan.result, nil
77}
78
79func (t *ASTTool) Execute(ctx context.Context, raw json.RawMessage, _ func(string)) (core.ToolResult, error) {
80 args, err := parseASTArgs(raw)
81 if err != nil {
82 return core.ToolResult{}, err
83 }
84 plan, err := t.plan(ctx, args)
85 if err != nil {
86 return core.ToolResult{}, err
87 }
88 if args.Rewrite == "" {
89 return plan.result, nil
90 }
91 if err := t.applyRewrite(plan.changes); err != nil {
92 return core.ToolResult{}, err
93 }
94 if t.LSPDiagnostics {
95 attachMutationDiagnostics(ctx, t.CWD, plan.result.Context.Mutates, t.LSP, &plan.result)
96 }
97 return plan.result, nil
98}
99
100func (t *ASTTool) applyRewrite(changes []astChange) error {
101 // Validate the entire plan before starting any writes, so a stale plan
102 // is rejected without touching any file.
103 for _, change := range changes {
104 if _, err := t.validateASTChange(change); err != nil {
105 return err
106 }
107 }
108 var written []string
109 for _, change := range changes {
110 // Validate again immediately before replacing each file: earlier
111 // writes take time, and a symlink may have been retargeted.
112 target, err := t.validateASTChange(change)
113 if err == nil {
114 err = writeASTChange(target, change)
115 }
116 if err != nil {
117 if len(written) == 0 {
118 return err
119 }
120 return fmt.Errorf("%w (already written: %s)", err, strings.Join(written, ", "))
121 }
122 written = append(written, change.path)
123 }
124 return nil
125}
126
127// validateASTChange resolves the file that will be replaced, checks it
128// against the write scope, and confirms it still holds the planned bytes.
129func (t *ASTTool) validateASTChange(change astChange) (string, error) {
130 if err := t.Sandbox.CheckWritePath(change.path); err != nil {
131 return "", err
132 }
133 // Rename replaces a symlink itself, so write through to its target like
134 // os.WriteFile does, and check that target too.
135 target, err := filepath.EvalSymlinks(change.path)
136 if err != nil {
137 return "", fmt.Errorf("ast: resolve %s: %w", change.path, err)
138 }
139 if err := t.Sandbox.CheckWritePath(target); err != nil {
140 return "", err
141 }
142 current, err := os.ReadFile(target)
143 if err != nil {
144 return "", fmt.Errorf("ast: re-read %s: %w", change.path, err)
145 }
146 if !bytes.Equal(current, change.original) {
147 return "", fmt.Errorf("ast: %s changed while applying rewrite", change.path)
148 }
149 return target, nil
150}
151
152type astChange struct {
153 path string
154 original []byte
155 content []byte
156 mode os.FileMode
157}
158
159// writeASTChange replaces target through a temporary file in the same
160// directory, so a failed write never leaves a truncated file.
161func writeASTChange(target string, change astChange) error {
162 tmp, err := os.CreateTemp(filepath.Dir(target), ".zut-ast-*")
163 if err != nil {
164 return fmt.Errorf("ast: write %s: %w", change.path, err)
165 }
166 defer os.Remove(tmp.Name())
167 err = tmp.Chmod(change.mode)
168 if err == nil {
169 _, err = tmp.Write(change.content)
170 }
171 if err == nil {
172 err = tmp.Sync()
173 }
174 if closeErr := tmp.Close(); err == nil {
175 err = closeErr
176 }
177 if err == nil {
178 err = os.Rename(tmp.Name(), target)
179 }
180 if err != nil {
181 return fmt.Errorf("ast: write %s: %w", change.path, err)
182 }
183 return nil
184}
185
186type astPlan struct {
187 result core.ToolResult
188 changes []astChange
189}
190
191type astMatch struct {
192 Text string `json:"text"`
193 File string `json:"file"`
194 Lines string `json:"lines"`
195 Replacement string `json:"replacement"`
196 Range struct {
197 ByteOffset struct {
198 Start int `json:"start"`
199 End int `json:"end"`
200 } `json:"byteOffset"`
201 Start struct {
202 Line int `json:"line"`
203 } `json:"start"`
204 } `json:"range"`
205}
206
207func parseASTArgs(raw json.RawMessage) (astArgs, error) {
208 var args astArgs
209 if err := json.Unmarshal(raw, &args); err != nil {
210 return args, fmt.Errorf("ast: invalid arguments: %w", err)
211 }
212 args.Pattern = strings.TrimSpace(args.Pattern)
213 args.Language = strings.TrimSpace(args.Language)
214 if args.Pattern == "" || strings.TrimSpace(args.Path) == "" {
215 return args, fmt.Errorf("ast: pattern and path are required")
216 }
217 return args, nil
218}
219
220func (t *ASTTool) plan(ctx context.Context, args astArgs) (astPlan, error) {
221 if t.ReadOnly && args.Rewrite != "" {
222 return astPlan{}, fmt.Errorf("ast: rewrites are unavailable in read-only orchestration")
223 }
224 root, err := canonicalPath(t.CWD, args.Path)
225 if err != nil {
226 return astPlan{}, fmt.Errorf("ast: resolve path: %w", err)
227 }
228 if args.Rewrite == "" {
229 if err := t.Sandbox.CheckReadPath(root); err != nil {
230 return astPlan{}, err
231 }
232 } else if err := t.Sandbox.CheckWritePath(root); err != nil {
233 return astPlan{}, err
234 }
235 info, err := os.Stat(root)
236 if err != nil {
237 return astPlan{}, fmt.Errorf("ast: %w", err)
238 }
239 if !info.IsDir() && !info.Mode().IsRegular() {
240 return astPlan{}, fmt.Errorf("ast: %s is not a regular file or directory", args.Path)
241 }
242 binary, err := t.executable()
243 if err != nil {
244 return astPlan{}, err
245 }
246 dir := root
247 if !info.IsDir() {
248 dir = filepath.Dir(root)
249 }
250 commandArgs := []string{"run", "--pattern=" + args.Pattern}
251 if args.Language != "" {
252 commandArgs = append(commandArgs, "--lang="+args.Language)
253 }
254 if args.Rewrite != "" {
255 commandArgs = append(commandArgs, "--rewrite="+args.Rewrite)
256 }
257 commandArgs = append(commandArgs, "--json=stream", root)
258 run, err := runASTGrep(ctx, binary, dir, commandArgs, args.Rewrite != "")
259 if err != nil {
260 return astPlan{}, err
261 }
262 if len(run.matches) == 0 && astGoCallPattern(args, root) {
263 // tree-sitter-go parses a bare `pkg.Func(...)` pattern as a type
264 // conversion, so it never matches real calls. Retry it as a call
265 // expression inside a function body. A zero-match retry keeps the
266 // original result and its warnings; a failed retry is reported.
267 retry, err := runASTGrep(ctx, binary, dir, astGoCallRuleArgs(args, root), args.Rewrite != "")
268 if err != nil {
269 return astPlan{}, fmt.Errorf("%w (while retrying the Go call pattern as a call expression)", err)
270 }
271 if len(retry.matches) > 0 {
272 run = retry
273 }
274 }
275 matches, stderr := run.matches, run.stderr
276 for index := range matches {
277 if !filepath.IsAbs(matches[index].File) {
278 matches[index].File = filepath.Join(dir, matches[index].File)
279 }
280 matches[index].File = filepath.Clean(matches[index].File)
281 }
282 if args.Rewrite == "" {
283 text, paths := formatASTMatches(matches)
284 if len(matches) == 0 {
285 text += astWarning(stderr) + astNoMatchHint
286 }
287 if run.overflow {
288 text += fmt.Sprintf("\n... [search output truncated at %d bytes; narrow path or pattern for complete results]", maxASTCaptureBytes)
289 }
290 return astPlan{result: core.ToolResult{
291 Content: []provider.Content{provider.TextBlock{Text: text}},
292 Context: provider.ToolContext{Discovers: paths},
293 Details: map[string]any{"engine": "ast-grep", "path": root, "language": args.Language, "matches": len(matches)},
294 }}, nil
295 }
296 plan, err := t.planRewrite(root, args, matches)
297 if err == nil && len(matches) == 0 {
298 text := plan.result.Content[0].(provider.TextBlock).Text + astWarning(stderr) + astNoMatchHint
299 plan.result.Content = []provider.Content{provider.TextBlock{Text: text}}
300 plan.result.Details.(map[string]any)["diff"] = text
301 }
302 return plan, err
303}
304
305// Available reports whether an ast-grep executable can be found.
306func (t *ASTTool) Available() bool {
307 _, err := t.executable()
308 return err == nil
309}
310
311func (t *ASTTool) executable() (string, error) {
312 lookup := t.LookPath
313 if lookup == nil {
314 lookup = exec.LookPath
315 }
316 if path, err := lookup("ast-grep"); err == nil && path != "" {
317 return path, nil
318 }
319 if path, err := lookup("sg"); err == nil && path != "" {
320 ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
321 defer cancel()
322 output, err := exec.CommandContext(ctx, path, "--version").CombinedOutput()
323 if err == nil && strings.Contains(strings.ToLower(string(output)), "ast-grep") {
324 return path, nil
325 }
326 }
327 return "", fmt.Errorf("ast: ast-grep executable ('ast-grep' or 'sg') not found in PATH")
328}
329
330const astNoMatchHint = "\nTip: the pattern must be valid code in the target language; $A matches one node and $$$A matches zero or more. Set language when the path mixes languages."
331
332type astRun struct {
333 matches []astMatch
334 stderr string
335 overflow bool
336}
337
338// runASTGrep runs one ast-grep command and decodes its JSON stream. Exit
339// status 1 with no output is ast-grep's normal "no matches" result.
340func runASTGrep(ctx context.Context, binary, dir string, commandArgs []string, rewrite bool) (astRun, error) {
341 cmd := exec.CommandContext(ctx, binary, commandArgs...)
342 cmd.Dir = dir
343 configureBashProcess(cmd, nil)
344 var stdout, stderr bytes.Buffer
345 capture := &limitedASTBuffer{buffer: &stdout, limit: maxASTCaptureBytes}
346 cmd.Stdout = capture
347 cmd.Stderr = &stderr
348 runErr := cmd.Run()
349 if ctx.Err() != nil {
350 return astRun{}, ctx.Err()
351 }
352 if capture.overflow {
353 if rewrite {
354 return astRun{}, fmt.Errorf("ast: output exceeds %d bytes; use a narrower path or pattern", capture.limit)
355 }
356 // The capture can end mid-record. Only decode complete JSON lines.
357 data := append([]byte(nil), stdout.Bytes()...)
358 stdout.Reset()
359 if end := bytes.LastIndexByte(data, '\n'); end >= 0 {
360 _, _ = stdout.Write(data[:end+1])
361 }
362 }
363 if runErr != nil {
364 var exit *exec.ExitError
365 if !errors.As(runErr, &exit) || exit.ExitCode() != 1 || strings.TrimSpace(stderr.String()) != "" || stdout.Len() != 0 {
366 message := strings.TrimSpace(stderr.String())
367 if message == "" {
368 message = runErr.Error()
369 }
370 return astRun{}, fmt.Errorf("ast: %s", message)
371 }
372 }
373 matches, err := decodeASTMatches(stdout.Bytes())
374 if err != nil {
375 return astRun{}, err
376 }
377 return astRun{matches: matches, stderr: stderr.String(), overflow: capture.overflow}, nil
378}
379
380var astGoQualifiedCall = regexp.MustCompile(`^[A-Za-z_$][\w$]*\.[A-Za-z_$][\w$]*\(.*\)$`)
381
382// astGoCallPattern reports whether a Go search may have hit the tree-sitter
383// ambiguity where `pkg.Func(x)` parses as a type conversion.
384func astGoCallPattern(args astArgs, root string) bool {
385 switch {
386 case strings.EqualFold(args.Language, "go"), strings.EqualFold(args.Language, "golang"):
387 case args.Language == "" && strings.EqualFold(filepath.Ext(root), ".go"):
388 case args.Language == "" && filepath.Ext(root) == "":
389 default:
390 return false
391 }
392 return astGoQualifiedCall.MatchString(args.Pattern)
393}
394
395// astGoCallRuleArgs matches the pattern as a call expression inside a
396// function body, which removes the type-conversion ambiguity.
397func astGoCallRuleArgs(args astArgs, root string) []string {
398 rule := map[string]any{
399 "id": "zut-ast",
400 "language": "go",
401 "rule": map[string]any{"pattern": map[string]any{
402 "context": "func _() { " + args.Pattern + " }",
403 "selector": "call_expression",
404 }},
405 }
406 if args.Rewrite != "" {
407 rule["fix"] = args.Rewrite
408 }
409 encoded, _ := json.Marshal(rule)
410 return []string{"scan", "--inline-rules=" + string(encoded), "--json=stream", root}
411}
412
413type limitedASTBuffer struct {
414 buffer *bytes.Buffer
415 limit int
416 overflow bool
417}
418
419func (w *limitedASTBuffer) Write(p []byte) (int, error) {
420 remaining := w.limit - w.buffer.Len()
421 if remaining > len(p) {
422 remaining = len(p)
423 }
424 if remaining > 0 {
425 _, _ = w.buffer.Write(p[:remaining])
426 }
427 if remaining < len(p) {
428 w.overflow = true
429 }
430 // Always consume the complete write so the child can exit normally. Returning
431 // an error here closes its stdout pipe and masks this limit as "broken pipe".
432 return len(p), nil
433}
434
435func decodeASTMatches(output []byte) ([]astMatch, error) {
436 var matches []astMatch
437 scanner := bufio.NewScanner(bytes.NewReader(output))
438 scanner.Buffer(make([]byte, 64*1024), 2*1024*1024)
439 for scanner.Scan() {
440 var match astMatch
441 if err := json.Unmarshal(scanner.Bytes(), &match); err != nil {
442 return nil, fmt.Errorf("ast: decode result: %w", err)
443 }
444 if match.File != "" {
445 matches = append(matches, match)
446 }
447 }
448 if err := scanner.Err(); err != nil {
449 return nil, fmt.Errorf("ast: decode output: %w", err)
450 }
451 return matches, nil
452}
453
454func formatASTMatches(matches []astMatch) (string, []string) {
455 seen := make(map[string]bool)
456 var paths []string
457 var output strings.Builder
458 for _, match := range matches {
459 path := filepath.Clean(match.File)
460 fmt.Fprintf(&output, "%s:%d: %s\n", path, match.Range.Start.Line+1, strings.TrimSpace(match.Lines))
461 if !seen[path] {
462 seen[path] = true
463 paths = append(paths, path)
464 }
465 }
466 if len(matches) == 0 {
467 return "No structural matches.", nil
468 }
469 return boundASTText(strings.TrimSpace(output.String())), paths
470}
471
472func astWarning(stderr string) string {
473 stderr = strings.TrimSpace(stderr)
474 if stderr == "" {
475 return ""
476 }
477 if len(stderr) > 2048 {
478 stderr = stderr[:2048] + "... [warning truncated]"
479 }
480 return "\n" + stderr
481}
482
483func (t *ASTTool) planRewrite(root string, args astArgs, matches []astMatch) (astPlan, error) {
484 byFile := make(map[string][]astMatch)
485 for _, match := range matches {
486 path, err := canonicalPath(t.CWD, match.File)
487 if err != nil {
488 return astPlan{}, err
489 }
490 if err := t.Sandbox.CheckWritePath(path); err != nil {
491 return astPlan{}, err
492 }
493 byFile[path] = append(byFile[path], match)
494 }
495 paths := make([]string, 0, len(byFile))
496 for path := range byFile {
497 paths = append(paths, path)
498 }
499 sort.Strings(paths)
500 var changes []astChange
501 var diffs strings.Builder
502 for _, path := range paths {
503 original, err := os.ReadFile(path)
504 if err != nil {
505 return astPlan{}, fmt.Errorf("ast: read %s: %w", path, err)
506 }
507 info, err := os.Stat(path)
508 if err != nil {
509 return astPlan{}, err
510 }
511 updated := append([]byte(nil), original...)
512 fileMatches := byFile[path]
513 sort.Slice(fileMatches, func(i, j int) bool {
514 left, right := fileMatches[i].Range.ByteOffset, fileMatches[j].Range.ByteOffset
515 if left.Start == right.Start {
516 return left.End > right.End
517 }
518 return left.Start < right.Start
519 })
520 selected := make([]astMatch, 0, len(fileMatches))
521 lastEnd := -1
522 for _, match := range fileMatches {
523 start, end := match.Range.ByteOffset.Start, match.Range.ByteOffset.End
524 if start < 0 || end < start || end > len(original) {
525 return astPlan{}, fmt.Errorf("ast: invalid rewrite range in %s", path)
526 }
527 if start < lastEnd {
528 continue // ast-grep applies the outermost non-overlapping match.
529 }
530 if !bytes.Equal(original[start:end], []byte(match.Text)) {
531 return astPlan{}, fmt.Errorf("ast: %s changed while planning rewrite", path)
532 }
533 selected = append(selected, match)
534 lastEnd = end
535 }
536 for index := len(selected) - 1; index >= 0; index-- {
537 match := selected[index]
538 start, end := match.Range.ByteOffset.Start, match.Range.ByteOffset.End
539 updated = append(updated[:start], append([]byte(match.Replacement), updated[end:]...)...)
540 }
541 if bytes.Equal(original, updated) {
542 continue
543 }
544 if strings.EqualFold(args.Language, "go") || strings.EqualFold(args.Language, "golang") || (args.Language == "" && strings.EqualFold(filepath.Ext(path), ".go")) {
545 crlf := bytes.Contains(original, []byte("\r\n"))
546 normalized := bytes.ReplaceAll(original, []byte("\r\n"), []byte("\n"))
547 // Mixed endings and pre-existing gofmt differences must not cause
548 // unrelated lines to change as a side effect of this rewrite.
549 if !crlf || bytes.Count(original, []byte("\n")) == bytes.Count(original, []byte("\r\n")) {
550 if clean, err := format.Source(normalized); err == nil && bytes.Equal(clean, normalized) {
551 formatted, err := format.Source(updated)
552 if err == nil {
553 if crlf {
554 formatted = bytes.ReplaceAll(formatted, []byte("\n"), []byte("\r\n"))
555 }
556 updated = formatted
557 }
558 }
559 }
560 }
561 if bytes.Equal(original, updated) {
562 continue
563 }
564 diffs.WriteString(unifiedDiff(path, string(original), string(updated)))
565 changes = append(changes, astChange{path: path, original: original, content: updated, mode: info.Mode().Perm()})
566 }
567 text := boundASTText(diffs.String())
568 if len(changes) == 0 {
569 if len(matches) == 0 {
570 text = "No structural matches."
571 } else {
572 text = "Structural matches found, but the rewrite produced no changes."
573 }
574 }
575 mutates := make([]string, 0, len(changes))
576 for _, change := range changes {
577 mutates = append(mutates, change.path)
578 }
579 return astPlan{changes: changes, result: core.ToolResult{
580 Content: []provider.Content{provider.TextBlock{Text: text}},
581 Context: provider.ToolContext{Mutates: mutates},
582 Details: map[string]any{"engine": "ast-grep", "path": root, "language": args.Language, "matches": len(matches), "files": len(changes), "diff": text},
583 }}, nil
584}
585
586func canonicalPath(cwd, path string) (string, error) {
587 path = resolvePath(cwd, path)
588 absolute, err := filepath.Abs(path)
589 if err != nil {
590 return "", err
591 }
592 return filepath.Clean(absolute), nil
593}
594
595func boundASTText(text string) string {
596 if len(text) <= maxASTOutputBytes {
597 return text
598 }
599 suffix := fmt.Sprintf("\n... [truncated at %d bytes]", maxASTOutputBytes)
600 text = text[:maxASTOutputBytes-len(suffix)]
601 for len(text) > 0 && !utf8.ValidString(text) {
602 text = text[:len(text)-1]
603 }
604 return text + suffix
605}