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 "sort"
15 "strings"
16 "unicode/utf8"
17
18 "github.com/bnema/zut/packages/agent/lsp"
19 "github.com/bnema/zut/packages/core"
20 "github.com/bnema/zut/packages/provider"
21)
22
23const (
24 maxASTOutputBytes = 60 * 1024
25 maxASTCaptureBytes = 4 * 1024 * 1024
26)
27
28// ASTTool performs structural searches and rewrites through ast-grep.
29type ASTTool struct {
30 CWD string
31 Sandbox *Sandbox
32 LSP *lsp.Manager
33 LSPDiagnostics bool
34 ReadOnly bool
35 LookPath func(string) (string, error)
36}
37
38var (
39 _ core.Tool = (*ASTTool)(nil)
40 _ core.ToolPreviewer = (*ASTTool)(nil)
41)
42
43type astArgs struct {
44 Pattern string `json:"pattern"`
45 Language string `json:"language"`
46 Path string `json:"path"`
47 Rewrite string `json:"rewrite,omitempty"`
48}
49
50const astSchema = `{"type":"object","properties":{"pattern":{"type":"string","description":"ast-grep structural pattern, such as '$A($$$ARGS)' for a call."},"language":{"type":"string","description":"ast-grep language id, such as go, rust, python, or typescript."},"path":{"type":"string","description":"File or directory to search, relative to the workspace or absolute."},"rewrite":{"type":"string","description":"Optional structural replacement. When supplied, Zut previews and safely applies every match."}},"required":["pattern","language","path"],"additionalProperties":false}`
51
52func (t *ASTTool) Name() string { return "ast" }
53func (t *ASTTool) Description() string {
54 return "Prefer this for syntax-aware code searches and rewrites. Uses installed ast-grep, previews mutations, enforces filesystem scope, and attaches diagnostics. Use grep/edit for plain text."
55}
56func (t *ASTTool) Schema() json.RawMessage { return json.RawMessage(astSchema) }
57
58func (t *ASTTool) Preview(ctx context.Context, raw json.RawMessage) (core.ToolResult, error) {
59 args, err := parseASTArgs(raw)
60 if err != nil {
61 return core.ToolResult{}, err
62 }
63 plan, err := t.plan(ctx, args)
64 if err != nil {
65 return core.ToolResult{}, err
66 }
67 return plan.result, nil
68}
69
70func (t *ASTTool) Execute(ctx context.Context, raw json.RawMessage, _ func(string)) (core.ToolResult, error) {
71 args, err := parseASTArgs(raw)
72 if err != nil {
73 return core.ToolResult{}, err
74 }
75 plan, err := t.plan(ctx, args)
76 if err != nil {
77 return core.ToolResult{}, err
78 }
79 if args.Rewrite == "" {
80 return plan.result, nil
81 }
82 for _, change := range plan.changes {
83 if err := os.WriteFile(change.path, change.content, change.mode); err != nil {
84 return core.ToolResult{}, fmt.Errorf("ast: write %s: %w", change.path, err)
85 }
86 }
87 if t.LSPDiagnostics {
88 attachMutationDiagnostics(ctx, t.CWD, plan.result.Context.Mutates, t.LSP, &plan.result)
89 }
90 return plan.result, nil
91}
92
93type astChange struct {
94 path string
95 content []byte
96 mode os.FileMode
97}
98
99type astPlan struct {
100 result core.ToolResult
101 changes []astChange
102}
103
104type astMatch struct {
105 Text string `json:"text"`
106 File string `json:"file"`
107 Lines string `json:"lines"`
108 Replacement string `json:"replacement"`
109 Range struct {
110 ByteOffset struct {
111 Start int `json:"start"`
112 End int `json:"end"`
113 } `json:"byteOffset"`
114 Start struct {
115 Line int `json:"line"`
116 } `json:"start"`
117 } `json:"range"`
118}
119
120func parseASTArgs(raw json.RawMessage) (astArgs, error) {
121 var args astArgs
122 if err := json.Unmarshal(raw, &args); err != nil {
123 return args, fmt.Errorf("ast: invalid arguments: %w", err)
124 }
125 args.Pattern = strings.TrimSpace(args.Pattern)
126 args.Language = strings.TrimSpace(args.Language)
127 if args.Pattern == "" || args.Language == "" || strings.TrimSpace(args.Path) == "" {
128 return args, fmt.Errorf("ast: pattern, language, and path are required")
129 }
130 return args, nil
131}
132
133func (t *ASTTool) plan(ctx context.Context, args astArgs) (astPlan, error) {
134 if t.ReadOnly && args.Rewrite != "" {
135 return astPlan{}, fmt.Errorf("ast: rewrites are unavailable in read-only orchestration")
136 }
137 root, err := canonicalPath(t.CWD, args.Path)
138 if err != nil {
139 return astPlan{}, fmt.Errorf("ast: resolve path: %w", err)
140 }
141 if args.Rewrite == "" {
142 if err := t.Sandbox.CheckReadPath(root); err != nil {
143 return astPlan{}, err
144 }
145 } else if err := t.Sandbox.CheckWritePath(root); err != nil {
146 return astPlan{}, err
147 }
148 info, err := os.Stat(root)
149 if err != nil {
150 return astPlan{}, fmt.Errorf("ast: %w", err)
151 }
152 if !info.IsDir() && !info.Mode().IsRegular() {
153 return astPlan{}, fmt.Errorf("ast: %s is not a regular file or directory", args.Path)
154 }
155 binary, err := t.executable()
156 if err != nil {
157 return astPlan{}, err
158 }
159 commandArgs := []string{"run", "--pattern=" + args.Pattern, "--lang=" + args.Language}
160 if args.Rewrite != "" {
161 commandArgs = append(commandArgs, "--rewrite="+args.Rewrite)
162 }
163 commandArgs = append(commandArgs, "--json=stream", root)
164 cmd := exec.CommandContext(ctx, binary, commandArgs...)
165 cmd.Dir = root
166 if !info.IsDir() {
167 cmd.Dir = filepath.Dir(root)
168 }
169 configureBashProcess(cmd, nil)
170 var stdout, stderr bytes.Buffer
171 capture := &limitedASTBuffer{buffer: &stdout, limit: maxASTCaptureBytes}
172 cmd.Stdout = capture
173 cmd.Stderr = &stderr
174 runErr := cmd.Run()
175 if ctx.Err() != nil {
176 return astPlan{}, ctx.Err()
177 }
178 if capture.overflow {
179 return astPlan{}, fmt.Errorf("ast: output exceeds %d bytes; use a narrower path or pattern", capture.limit)
180 }
181 if runErr != nil {
182 var exit *exec.ExitError
183 if !errors.As(runErr, &exit) || exit.ExitCode() != 1 || strings.TrimSpace(stderr.String()) != "" || stdout.Len() != 0 {
184 message := strings.TrimSpace(stderr.String())
185 if message == "" {
186 message = runErr.Error()
187 }
188 return astPlan{}, fmt.Errorf("ast: %s", message)
189 }
190 }
191 matches, err := decodeASTMatches(stdout.Bytes())
192 if err != nil {
193 return astPlan{}, err
194 }
195 for index := range matches {
196 if !filepath.IsAbs(matches[index].File) {
197 matches[index].File = filepath.Join(cmd.Dir, matches[index].File)
198 }
199 matches[index].File = filepath.Clean(matches[index].File)
200 }
201 if args.Rewrite == "" {
202 text, paths := formatASTMatches(matches)
203 return astPlan{result: core.ToolResult{
204 Content: []provider.Content{provider.TextBlock{Text: text}},
205 Context: provider.ToolContext{Discovers: paths},
206 Details: map[string]any{"engine": "ast-grep", "path": root, "language": args.Language, "matches": len(matches)},
207 }}, nil
208 }
209 return t.planRewrite(root, args, matches)
210}
211
212func (t *ASTTool) executable() (string, error) {
213 lookup := t.LookPath
214 if lookup == nil {
215 lookup = exec.LookPath
216 }
217 for _, name := range []string{"ast-grep", "sg"} {
218 if path, err := lookup(name); err == nil && path != "" {
219 return path, nil
220 }
221 }
222 return "", fmt.Errorf("ast: ast-grep executable ('ast-grep' or 'sg') not found in PATH")
223}
224
225type limitedASTBuffer struct {
226 buffer *bytes.Buffer
227 limit int
228 overflow bool
229}
230
231func (w *limitedASTBuffer) Write(p []byte) (int, error) {
232 remaining := w.limit - w.buffer.Len()
233 if remaining > len(p) {
234 remaining = len(p)
235 }
236 if remaining > 0 {
237 _, _ = w.buffer.Write(p[:remaining])
238 }
239 if remaining < len(p) {
240 w.overflow = true
241 }
242 // Always consume the complete write so the child can exit normally. Returning
243 // an error here closes its stdout pipe and masks this limit as "broken pipe".
244 return len(p), nil
245}
246
247func decodeASTMatches(output []byte) ([]astMatch, error) {
248 var matches []astMatch
249 scanner := bufio.NewScanner(bytes.NewReader(output))
250 scanner.Buffer(make([]byte, 64*1024), 2*1024*1024)
251 for scanner.Scan() {
252 var match astMatch
253 if err := json.Unmarshal(scanner.Bytes(), &match); err != nil {
254 return nil, fmt.Errorf("ast: decode result: %w", err)
255 }
256 if match.File != "" {
257 matches = append(matches, match)
258 }
259 }
260 if err := scanner.Err(); err != nil {
261 return nil, fmt.Errorf("ast: decode output: %w", err)
262 }
263 return matches, nil
264}
265
266func formatASTMatches(matches []astMatch) (string, []string) {
267 seen := make(map[string]bool)
268 var paths []string
269 var output strings.Builder
270 for _, match := range matches {
271 path := filepath.Clean(match.File)
272 fmt.Fprintf(&output, "%s:%d: %s\n", path, match.Range.Start.Line+1, strings.TrimSpace(match.Lines))
273 if !seen[path] {
274 seen[path] = true
275 paths = append(paths, path)
276 }
277 }
278 if len(matches) == 0 {
279 return "No structural matches.", nil
280 }
281 return boundASTText(strings.TrimSpace(output.String())), paths
282}
283
284func (t *ASTTool) planRewrite(root string, args astArgs, matches []astMatch) (astPlan, error) {
285 byFile := make(map[string][]astMatch)
286 for _, match := range matches {
287 path, err := canonicalPath(t.CWD, match.File)
288 if err != nil {
289 return astPlan{}, err
290 }
291 if err := t.Sandbox.CheckWritePath(path); err != nil {
292 return astPlan{}, err
293 }
294 byFile[path] = append(byFile[path], match)
295 }
296 paths := make([]string, 0, len(byFile))
297 for path := range byFile {
298 paths = append(paths, path)
299 }
300 sort.Strings(paths)
301 var changes []astChange
302 var diffs strings.Builder
303 for _, path := range paths {
304 original, err := os.ReadFile(path)
305 if err != nil {
306 return astPlan{}, fmt.Errorf("ast: read %s: %w", path, err)
307 }
308 info, err := os.Stat(path)
309 if err != nil {
310 return astPlan{}, err
311 }
312 updated := append([]byte(nil), original...)
313 fileMatches := byFile[path]
314 sort.Slice(fileMatches, func(i, j int) bool {
315 left, right := fileMatches[i].Range.ByteOffset, fileMatches[j].Range.ByteOffset
316 if left.Start == right.Start {
317 return left.End > right.End
318 }
319 return left.Start < right.Start
320 })
321 selected := make([]astMatch, 0, len(fileMatches))
322 lastEnd := -1
323 for _, match := range fileMatches {
324 start, end := match.Range.ByteOffset.Start, match.Range.ByteOffset.End
325 if start < 0 || end < start || end > len(original) {
326 return astPlan{}, fmt.Errorf("ast: invalid rewrite range in %s", path)
327 }
328 if start < lastEnd {
329 continue // ast-grep applies the outermost non-overlapping match.
330 }
331 if !bytes.Equal(original[start:end], []byte(match.Text)) {
332 return astPlan{}, fmt.Errorf("ast: %s changed while planning rewrite", path)
333 }
334 selected = append(selected, match)
335 lastEnd = end
336 }
337 for index := len(selected) - 1; index >= 0; index-- {
338 match := selected[index]
339 start, end := match.Range.ByteOffset.Start, match.Range.ByteOffset.End
340 updated = append(updated[:start], append([]byte(match.Replacement), updated[end:]...)...)
341 }
342 if bytes.Equal(original, updated) {
343 continue
344 }
345 if strings.EqualFold(args.Language, "go") {
346 formatted, formatErr := format.Source(updated)
347 if formatErr == nil {
348 if bytes.Contains(original, []byte("\r\n")) {
349 formatted = bytes.ReplaceAll(formatted, []byte("\n"), []byte("\r\n"))
350 }
351 updated = formatted
352 }
353 }
354 if bytes.Equal(original, updated) {
355 continue
356 }
357 diffs.WriteString(unifiedDiff(path, string(original), string(updated)))
358 changes = append(changes, astChange{path: path, content: updated, mode: info.Mode().Perm()})
359 }
360 text := boundASTText(diffs.String())
361 if len(changes) == 0 {
362 if len(matches) == 0 {
363 text = "No structural matches."
364 } else {
365 text = "Structural matches found, but the rewrite produced no changes."
366 }
367 }
368 mutates := make([]string, 0, len(changes))
369 for _, change := range changes {
370 mutates = append(mutates, change.path)
371 }
372 return astPlan{changes: changes, result: core.ToolResult{
373 Content: []provider.Content{provider.TextBlock{Text: text}},
374 Context: provider.ToolContext{Mutates: mutates},
375 Details: map[string]any{"engine": "ast-grep", "path": root, "language": args.Language, "matches": len(matches), "files": len(changes), "diff": text},
376 }}, nil
377}
378
379func canonicalPath(cwd, path string) (string, error) {
380 path = resolvePath(cwd, path)
381 absolute, err := filepath.Abs(path)
382 if err != nil {
383 return "", err
384 }
385 return filepath.Clean(absolute), nil
386}
387
388func boundASTText(text string) string {
389 if len(text) <= maxASTOutputBytes {
390 return text
391 }
392 suffix := fmt.Sprintf("\n... [truncated at %d bytes]", maxASTOutputBytes)
393 text = text[:maxASTOutputBytes-len(suffix)]
394 for len(text) > 0 && !utf8.ValidString(text) {
395 text = text[:len(text)-1]
396 }
397 return text + suffix
398}