feat(ast): harden rewrites and help agents adopt structural search (#232)
* feat(ast): harden rewrites and make structural patterns easier to use Verify an sg fallback is ast-grep, re-check files before writing, write via temp file and rename (through symlinks), keep truncated search results, surface ast-grep warnings, and limit gofmt to already-formatted files with consistent line endings. Make language optional, add usage examples to the tool description, add a no-match syntax tip, and retry Go qualified calls such as fmt.Errorf($$$) as call expressions since tree-sitter-go parses them as type conversions. * feat(grep): suggest ast for call and definition searches When grep receives a regex that looks like a call or definition search and the ast tool is registered with ast-grep installed, append a short ast example to the result. * fix(ast): revalidate symlink targets per write and report failed call retries Check each rewrite target's resolved path against the write scope and its current contents immediately before replacing it, report a failed Go call-expression retry instead of returning zero matches, and only suggest ast from grep for calls with an escaped dot. * test(ast): make path and permission assertions portable to Windows Build JSON arguments with json.Marshal so Windows backslash paths are valid escapes, and skip Unix permission-bit assertions on Windows.
- commit
- a3faff4be41f5390a76b120765e777393416e197
- author
- bnema · · 2026-09-29 04:12 UTC
- committer
- GitHub · · 2026-09-29 04:12 UTC
- parent
- cb619f7
- signature
- GPG Unverified
B5690EEEBB952194The key is unknown
7 files · +729 -60
M README.md +1 -1
| @@ -266,7 +266,7 @@ Print-mode stats contain `provider`, `model`, `prompt_tokens`, `reasoning_tokens | |||
| 266 | 266 | - `edit`: one or more exact-match replacements in an existing file. | |
| 267 | 267 | - `grep`: search a file or directory for a regular expression using ripgrep when available, with a portable grep fallback. Search output is bounded and follows the active filesystem read scope; it never invokes a shell. The fallback's regular-expression syntax and ignore/hidden-file behavior can differ from ripgrep. | |
| 268 | 268 | - `glob`: find files matching a glob pattern (`**/*.go`, `*.json`, `src/**/*.ts`) within a directory tree. Supports `*`, `?`, `**`, and character classes; patterns with a slash match relative paths, others match basenames. Honors `.gitignore` rules including nested files, skips hidden entries unless asked, never follows nested directory symlinks, and bounds depth (24) and matches (500, sorted). | |
| 269 | - | - `ast`: search or rewrite source code structurally with an [ast-grep](https://ast-grep.github.io/) pattern. Prefer it when the operation depends on code syntax; use `grep`/`edit` for plain text. Zut starts an installed `ast-grep` or `sg` executable directly without a shell. Searches return compact `path:line` matches. Rewrites are planned in memory, shown as a diff for confirmation, checked against the active filesystem scope, applied with file permissions preserved, and followed by bounded diagnostics. Install ast-grep separately; the tool returns a clear error when neither executable is available. | |
| 269 | + | - `ast`: search or rewrite source code structurally with an [ast-grep](https://ast-grep.github.io/) pattern. Prefer it when the operation depends on code syntax; use `grep`/`edit` for plain text. `language` is optional; when omitted, ast-grep infers it from each file's extension. Zut starts an installed `ast-grep` executable, or `sg` when it identifies itself as ast-grep, directly without a shell. Searches return compact `path:line` matches. Rewrites are planned in memory, shown as a diff for confirmation, checked against the active filesystem scope, verified against the files' current contents, written through a temporary file and rename with file permissions preserved, and followed by bounded diagnostics. When `grep` is given a regex that looks like a call or definition search and `ast` is available, its result includes a short `ast` suggestion. Install ast-grep separately; the tool returns a clear error when it is not available. | |
| 270 | 270 | - `bash`: run a command in the session cwd with merged stdout/stderr. Every call must include a positive integer `timeout` in seconds; calls without it or with `timeout <= 0` are rejected, and the deadline is always enforced. On Unix, zut uses `/bin/bash -c` when available, then `bash -c` from `PATH`, and falls back to POSIX `/bin/sh -c` when Bash is unavailable. On Windows, it uses `cmd /C`. macOS ships Bash 3.2 by default, so newer Bash features may be unavailable. | |
| 271 | 271 | - `python`: execute inline Python 3 code with an existing interpreter, never installing one. Every call must include nonblank `code` and a positive integer `timeout` in seconds (shared by interpreter discovery and execution). Each call runs a fresh stateless subprocess directly (`-u -B -`, code on stdin) in the session cwd with merged stdout/stderr; there is no persistent interpreter state, no script file, and no bytecode cache. Interactive input is unsupported (stdin observes EOF). Output is bounded (50 KiB, 2,000 lines) with explicit truncation marks; truncated results persist the byte-capped buffer to a temp file referenced as `full output` with its path in `full_output_path`, mirroring `bash`. Discovery order: `python_interpreter` in `$ZUT_HOME/config.json` (one executable path; relative paths resolve against the session cwd), then the active `VIRTUAL_ENV`, then `./.venv`, then `python3`/`python` on `PATH`. On Windows, Store aliases, `py` launchers, and install-manager entrypoints are rejected; only positively identified installed interpreters run. Missing packages fail with the ordinary import traceback. Tool arguments and output are recorded in the session like other tools; files the code deliberately writes persist. Save reusable scripts with `write`/`edit` instead of using `python` as an opaque file-edit path. Unavailable to packaged agents, and rejected while jailed (see below). | |
| 272 | 272 | - `worktree`: manage Git worktrees with the `create`, `list`, and `cleanup` actions. Creation checks out a new branch under `$ZUT_HOME/worktrees/<repo-id>`; an absolute local Git setting, `zut.worktrees.path`, overrides that root per repository. Listing inventories registered worktrees. Cleanup requires exact paths, confirmation through the normal preview gate, and revalidation immediately before removal. It refuses the primary worktree, dirty or detached checkouts, and commits not merged into the current branch. It removes neither local branches nor unmerged work. | |
| @@ -266,7 +266,7 @@ Print-mode stats contain `provider`, `model`, `prompt_tokens`, `reasoning_tokens | |||||
| 266 | - `edit`: one or more exact-match replacements in an existing file. | 266 | - `edit`: one or more exact-match replacements in an existing file. | ||
| 267 | - `grep`: search a file or directory for a regular expression using ripgrep when available, with a portable grep fallback. Search output is bounded and follows the active filesystem read scope; it never invokes a shell. The fallback's regular-expression syntax and ignore/hidden-file behavior can differ from ripgrep. | 267 | - `grep`: search a file or directory for a regular expression using ripgrep when available, with a portable grep fallback. Search output is bounded and follows the active filesystem read scope; it never invokes a shell. The fallback's regular-expression syntax and ignore/hidden-file behavior can differ from ripgrep. | ||
| 268 | - `glob`: find files matching a glob pattern (`**/*.go`, `*.json`, `src/**/*.ts`) within a directory tree. Supports `*`, `?`, `**`, and character classes; patterns with a slash match relative paths, others match basenames. Honors `.gitignore` rules including nested files, skips hidden entries unless asked, never follows nested directory symlinks, and bounds depth (24) and matches (500, sorted). | 268 | - `glob`: find files matching a glob pattern (`**/*.go`, `*.json`, `src/**/*.ts`) within a directory tree. Supports `*`, `?`, `**`, and character classes; patterns with a slash match relative paths, others match basenames. Honors `.gitignore` rules including nested files, skips hidden entries unless asked, never follows nested directory symlinks, and bounds depth (24) and matches (500, sorted). | ||
| 269 | - | - `ast`: search or rewrite source code structurally with an [ast-grep](https://ast-grep.github.io/) pattern. Prefer it when the operation depends on code syntax; use `grep`/`edit` for plain text. Zut starts an installed `ast-grep` or `sg` executable directly without a shell. Searches return compact `path:line` matches. Rewrites are planned in memory, shown as a diff for confirmation, checked against the active filesystem scope, applied with file permissions preserved, and followed by bounded diagnostics. Install ast-grep separately; the tool returns a clear error when neither executable is available. | 269 | + | - `ast`: search or rewrite source code structurally with an [ast-grep](https://ast-grep.github.io/) pattern. Prefer it when the operation depends on code syntax; use `grep`/`edit` for plain text. `language` is optional; when omitted, ast-grep infers it from each file's extension. Zut starts an installed `ast-grep` executable, or `sg` when it identifies itself as ast-grep, directly without a shell. Searches return compact `path:line` matches. Rewrites are planned in memory, shown as a diff for confirmation, checked against the active filesystem scope, verified against the files' current contents, written through a temporary file and rename with file permissions preserved, and followed by bounded diagnostics. When `grep` is given a regex that looks like a call or definition search and `ast` is available, its result includes a short `ast` suggestion. Install ast-grep separately; the tool returns a clear error when it is not available. |
| 270 | - `bash`: run a command in the session cwd with merged stdout/stderr. Every call must include a positive integer `timeout` in seconds; calls without it or with `timeout <= 0` are rejected, and the deadline is always enforced. On Unix, zut uses `/bin/bash -c` when available, then `bash -c` from `PATH`, and falls back to POSIX `/bin/sh -c` when Bash is unavailable. On Windows, it uses `cmd /C`. macOS ships Bash 3.2 by default, so newer Bash features may be unavailable. | 270 | - `bash`: run a command in the session cwd with merged stdout/stderr. Every call must include a positive integer `timeout` in seconds; calls without it or with `timeout <= 0` are rejected, and the deadline is always enforced. On Unix, zut uses `/bin/bash -c` when available, then `bash -c` from `PATH`, and falls back to POSIX `/bin/sh -c` when Bash is unavailable. On Windows, it uses `cmd /C`. macOS ships Bash 3.2 by default, so newer Bash features may be unavailable. | ||
| 271 | - `python`: execute inline Python 3 code with an existing interpreter, never installing one. Every call must include nonblank `code` and a positive integer `timeout` in seconds (shared by interpreter discovery and execution). Each call runs a fresh stateless subprocess directly (`-u -B -`, code on stdin) in the session cwd with merged stdout/stderr; there is no persistent interpreter state, no script file, and no bytecode cache. Interactive input is unsupported (stdin observes EOF). Output is bounded (50 KiB, 2,000 lines) with explicit truncation marks; truncated results persist the byte-capped buffer to a temp file referenced as `full output` with its path in `full_output_path`, mirroring `bash`. Discovery order: `python_interpreter` in `$ZUT_HOME/config.json` (one executable path; relative paths resolve against the session cwd), then the active `VIRTUAL_ENV`, then `./.venv`, then `python3`/`python` on `PATH`. On Windows, Store aliases, `py` launchers, and install-manager entrypoints are rejected; only positively identified installed interpreters run. Missing packages fail with the ordinary import traceback. Tool arguments and output are recorded in the session like other tools; files the code deliberately writes persist. Save reusable scripts with `write`/`edit` instead of using `python` as an opaque file-edit path. Unavailable to packaged agents, and rejected while jailed (see below). | 271 | - `python`: execute inline Python 3 code with an existing interpreter, never installing one. Every call must include nonblank `code` and a positive integer `timeout` in seconds (shared by interpreter discovery and execution). Each call runs a fresh stateless subprocess directly (`-u -B -`, code on stdin) in the session cwd with merged stdout/stderr; there is no persistent interpreter state, no script file, and no bytecode cache. Interactive input is unsupported (stdin observes EOF). Output is bounded (50 KiB, 2,000 lines) with explicit truncation marks; truncated results persist the byte-capped buffer to a temp file referenced as `full output` with its path in `full_output_path`, mirroring `bash`. Discovery order: `python_interpreter` in `$ZUT_HOME/config.json` (one executable path; relative paths resolve against the session cwd), then the active `VIRTUAL_ENV`, then `./.venv`, then `python3`/`python` on `PATH`. On Windows, Store aliases, `py` launchers, and install-manager entrypoints are rejected; only positively identified installed interpreters run. Missing packages fail with the ordinary import traceback. Tool arguments and output are recorded in the session like other tools; files the code deliberately writes persist. Save reusable scripts with `write`/`edit` instead of using `python` as an opaque file-edit path. Unavailable to packaged agents, and rejected while jailed (see below). | ||
| 272 | - `worktree`: manage Git worktrees with the `create`, `list`, and `cleanup` actions. Creation checks out a new branch under `$ZUT_HOME/worktrees/<repo-id>`; an absolute local Git setting, `zut.worktrees.path`, overrides that root per repository. Listing inventories registered worktrees. Cleanup requires exact paths, confirmation through the normal preview gate, and revalidation immediately before removal. It refuses the primary worktree, dirty or detached checkouts, and commits not merged into the current branch. It removes neither local branches nor unmerged work. | 272 | - `worktree`: manage Git worktrees with the `create`, `list`, and `cleanup` actions. Creation checks out a new branch under `$ZUT_HOME/worktrees/<repo-id>`; an absolute local Git setting, `zut.worktrees.path`, overrides that root per repository. Listing inventories registered worktrees. Cleanup requires exact paths, confirmation through the normal preview gate, and revalidation immediately before removal. It refuses the primary worktree, dirty or detached checkouts, and commits not merged into the current branch. It removes neither local branches nor unmerged work. | ||
M packages/agent/build.go +6 -0
| @@ -8,6 +8,7 @@ import ( | |||
| 8 | 8 | "runtime" | |
| 9 | 9 | "sort" | |
| 10 | 10 | "strings" | |
| 11 | + | "sync" | |
| 11 | 12 | ||
| 12 | 13 | zutdocs "github.com/bnema/zut" | |
| 13 | 14 | "github.com/bnema/zut/packages/agent/lsp" | |
| @@ -1436,6 +1437,11 @@ func buildToolRegistry(args Args, cwd string, sandbox *tools.Sandbox, lspEnabled | |||
| 1436 | 1437 | } | |
| 1437 | 1438 | } | |
| 1438 | 1439 | } | |
| 1440 | + | grepTool, hasGrep := reg["grep"].(*tools.GrepTool) | |
| 1441 | + | astTool, hasAST := reg["ast"].(*tools.ASTTool) | |
| 1442 | + | if hasGrep && hasAST { | |
| 1443 | + | grepTool.ASTAvailable = sync.OnceValue(astTool.Available) | |
| 1444 | + | } | |
| 1439 | 1445 | return reg | |
| 1440 | 1446 | } | |
| 1441 | 1447 | ||
| @@ -8,6 +8,7 @@ import ( | |||||
| 8 | "runtime" | 8 | "runtime" | ||
| 9 | "sort" | 9 | "sort" | ||
| 10 | "strings" | 10 | "strings" | ||
| 11 | + | "sync" | |||
| 11 | 12 | ||||
| 12 | zutdocs "github.com/bnema/zut" | 13 | zutdocs "github.com/bnema/zut" | ||
| 13 | "github.com/bnema/zut/packages/agent/lsp" | 14 | "github.com/bnema/zut/packages/agent/lsp" | ||
| @@ -1436,6 +1437,11 @@ func buildToolRegistry(args Args, cwd string, sandbox *tools.Sandbox, lspEnabled | |||||
| 1436 | } | 1437 | } | ||
| 1437 | } | 1438 | } | ||
| 1438 | } | 1439 | } | ||
| 1440 | + | grepTool, hasGrep := reg["grep"].(*tools.GrepTool) | |||
| 1441 | + | astTool, hasAST := reg["ast"].(*tools.ASTTool) | |||
| 1442 | + | if hasGrep && hasAST { | |||
| 1443 | + | grepTool.ASTAvailable = sync.OnceValue(astTool.Available) | |||
| 1444 | + | } | |||
| 1439 | return reg | 1445 | return reg | ||
| 1440 | } | 1446 | } | ||
| 1441 | 1447 | ||||
M packages/agent/tools/ast.go +258 -51
| @@ -11,8 +11,10 @@ import ( | |||
| 11 | 11 | "os" | |
| 12 | 12 | "os/exec" | |
| 13 | 13 | "path/filepath" | |
| 14 | + | "regexp" | |
| 14 | 15 | "sort" | |
| 15 | 16 | "strings" | |
| 17 | + | "time" | |
| 16 | 18 | "unicode/utf8" | |
| 17 | 19 | ||
| 18 | 20 | "github.com/bnema/zut/packages/agent/lsp" | |
| @@ -47,11 +49,18 @@ type astArgs struct { | |||
| 47 | 49 | Rewrite string `json:"rewrite,omitempty"` | |
| 48 | 50 | } | |
| 49 | 51 | ||
| 50 | - | const 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}` | |
| 52 | + | const 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}` | |
| 51 | 53 | ||
| 52 | 54 | func (t *ASTTool) Name() string { return "ast" } | |
| 53 | 55 | func (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." | |
| 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." | |
| 55 | 64 | } | |
| 56 | 65 | func (t *ASTTool) Schema() json.RawMessage { return json.RawMessage(astSchema) } | |
| 57 | 66 | ||
| @@ -79,10 +88,8 @@ func (t *ASTTool) Execute(ctx context.Context, raw json.RawMessage, _ func(strin | |||
| 79 | 88 | if args.Rewrite == "" { | |
| 80 | 89 | return plan.result, nil | |
| 81 | 90 | } | |
| 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 | - | } | |
| 91 | + | if err := t.applyRewrite(plan.changes); err != nil { | |
| 92 | + | return core.ToolResult{}, err | |
| 86 | 93 | } | |
| 87 | 94 | if t.LSPDiagnostics { | |
| 88 | 95 | attachMutationDiagnostics(ctx, t.CWD, plan.result.Context.Mutates, t.LSP, &plan.result) | |
| @@ -90,10 +97,90 @@ func (t *ASTTool) Execute(ctx context.Context, raw json.RawMessage, _ func(strin | |||
| 90 | 97 | return plan.result, nil | |
| 91 | 98 | } | |
| 92 | 99 | ||
| 100 | + | func (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. | |
| 129 | + | func (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 | + | ||
| 93 | 152 | type astChange struct { | |
| 94 | - | path string | |
| 95 | - | content []byte | |
| 96 | - | mode os.FileMode | |
| 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. | |
| 161 | + | func 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 | |
| 97 | 184 | } | |
| 98 | 185 | ||
| 99 | 186 | type astPlan struct { | |
| @@ -124,8 +211,8 @@ func parseASTArgs(raw json.RawMessage) (astArgs, error) { | |||
| 124 | 211 | } | |
| 125 | 212 | args.Pattern = strings.TrimSpace(args.Pattern) | |
| 126 | 213 | 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") | |
| 214 | + | if args.Pattern == "" || strings.TrimSpace(args.Path) == "" { | |
| 215 | + | return args, fmt.Errorf("ast: pattern and path are required") | |
| 129 | 216 | } | |
| 130 | 217 | return args, nil | |
| 131 | 218 | } | |
| @@ -156,57 +243,69 @@ func (t *ASTTool) plan(ctx context.Context, args astArgs) (astPlan, error) { | |||
| 156 | 243 | if err != nil { | |
| 157 | 244 | return astPlan{}, err | |
| 158 | 245 | } | |
| 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 | |
| 246 | + | dir := root | |
| 166 | 247 | 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() | |
| 248 | + | dir = filepath.Dir(root) | |
| 177 | 249 | } | |
| 178 | - | if capture.overflow { | |
| 179 | - | return astPlan{}, fmt.Errorf("ast: output exceeds %d bytes; use a narrower path or pattern", capture.limit) | |
| 250 | + | commandArgs := []string{"run", "--pattern=" + args.Pattern} | |
| 251 | + | if args.Language != "" { | |
| 252 | + | commandArgs = append(commandArgs, "--lang="+args.Language) | |
| 180 | 253 | } | |
| 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 | - | } | |
| 254 | + | if args.Rewrite != "" { | |
| 255 | + | commandArgs = append(commandArgs, "--rewrite="+args.Rewrite) | |
| 190 | 256 | } | |
| 191 | - | matches, err := decodeASTMatches(stdout.Bytes()) | |
| 257 | + | commandArgs = append(commandArgs, "--json=stream", root) | |
| 258 | + | run, err := runASTGrep(ctx, binary, dir, commandArgs, args.Rewrite != "") | |
| 192 | 259 | if err != nil { | |
| 193 | 260 | return astPlan{}, err | |
| 194 | 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 | |
| 195 | 276 | for index := range matches { | |
| 196 | 277 | if !filepath.IsAbs(matches[index].File) { | |
| 197 | - | matches[index].File = filepath.Join(cmd.Dir, matches[index].File) | |
| 278 | + | matches[index].File = filepath.Join(dir, matches[index].File) | |
| 198 | 279 | } | |
| 199 | 280 | matches[index].File = filepath.Clean(matches[index].File) | |
| 200 | 281 | } | |
| 201 | 282 | if args.Rewrite == "" { | |
| 202 | 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 | + | } | |
| 203 | 290 | return astPlan{result: core.ToolResult{ | |
| 204 | 291 | Content: []provider.Content{provider.TextBlock{Text: text}}, | |
| 205 | 292 | Context: provider.ToolContext{Discovers: paths}, | |
| 206 | 293 | Details: map[string]any{"engine": "ast-grep", "path": root, "language": args.Language, "matches": len(matches)}, | |
| 207 | 294 | }}, nil | |
| 208 | 295 | } | |
| 209 | - | return t.planRewrite(root, args, matches) | |
| 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. | |
| 306 | + | func (t *ASTTool) Available() bool { | |
| 307 | + | _, err := t.executable() | |
| 308 | + | return err == nil | |
| 210 | 309 | } | |
| 211 | 310 | ||
| 212 | 311 | func (t *ASTTool) executable() (string, error) { | |
| @@ -214,14 +313,103 @@ func (t *ASTTool) executable() (string, error) { | |||
| 214 | 313 | if lookup == nil { | |
| 215 | 314 | lookup = exec.LookPath | |
| 216 | 315 | } | |
| 217 | - | for _, name := range []string{"ast-grep", "sg"} { | |
| 218 | - | if path, err := lookup(name); err == nil && path != "" { | |
| 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") { | |
| 219 | 324 | return path, nil | |
| 220 | 325 | } | |
| 221 | 326 | } | |
| 222 | 327 | return "", fmt.Errorf("ast: ast-grep executable ('ast-grep' or 'sg') not found in PATH") | |
| 223 | 328 | } | |
| 224 | 329 | ||
| 330 | + | const 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 | + | ||
| 332 | + | type 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. | |
| 340 | + | func 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 | + | ||
| 380 | + | var 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. | |
| 384 | + | func 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. | |
| 397 | + | func 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 | + | ||
| 225 | 413 | type limitedASTBuffer struct { | |
| 226 | 414 | buffer *bytes.Buffer | |
| 227 | 415 | limit int | |
| @@ -281,6 +469,17 @@ func formatASTMatches(matches []astMatch) (string, []string) { | |||
| 281 | 469 | return boundASTText(strings.TrimSpace(output.String())), paths | |
| 282 | 470 | } | |
| 283 | 471 | ||
| 472 | + | func 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 | + | ||
| 284 | 483 | func (t *ASTTool) planRewrite(root string, args astArgs, matches []astMatch) (astPlan, error) { | |
| 285 | 484 | byFile := make(map[string][]astMatch) | |
| 286 | 485 | for _, match := range matches { | |
| @@ -342,20 +541,28 @@ func (t *ASTTool) planRewrite(root string, args astArgs, matches []astMatch) (as | |||
| 342 | 541 | if bytes.Equal(original, updated) { | |
| 343 | 542 | continue | |
| 344 | 543 | } | |
| 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")) | |
| 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 | + | } | |
| 350 | 558 | } | |
| 351 | - | updated = formatted | |
| 352 | 559 | } | |
| 353 | 560 | } | |
| 354 | 561 | if bytes.Equal(original, updated) { | |
| 355 | 562 | continue | |
| 356 | 563 | } | |
| 357 | 564 | diffs.WriteString(unifiedDiff(path, string(original), string(updated))) | |
| 358 | - | changes = append(changes, astChange{path: path, content: updated, mode: info.Mode().Perm()}) | |
| 565 | + | changes = append(changes, astChange{path: path, original: original, content: updated, mode: info.Mode().Perm()}) | |
| 359 | 566 | } | |
| 360 | 567 | text := boundASTText(diffs.String()) | |
| 361 | 568 | if len(changes) == 0 { | |
| @@ -11,8 +11,10 @@ import ( | |||||
| 11 | "os" | 11 | "os" | ||
| 12 | "os/exec" | 12 | "os/exec" | ||
| 13 | "path/filepath" | 13 | "path/filepath" | ||
| 14 | + | "regexp" | |||
| 14 | "sort" | 15 | "sort" | ||
| 15 | "strings" | 16 | "strings" | ||
| 17 | + | "time" | |||
| 16 | "unicode/utf8" | 18 | "unicode/utf8" | ||
| 17 | 19 | ||||
| 18 | "github.com/bnema/zut/packages/agent/lsp" | 20 | "github.com/bnema/zut/packages/agent/lsp" | ||
| @@ -47,11 +49,18 @@ type astArgs struct { | |||||
| 47 | Rewrite string `json:"rewrite,omitempty"` | 49 | Rewrite string `json:"rewrite,omitempty"` | ||
| 48 | } | 50 | } | ||
| 49 | 51 | ||||
| 50 | - | const 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}` | 52 | + | const 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}` |
| 51 | 53 | ||||
| 52 | func (t *ASTTool) Name() string { return "ast" } | 54 | func (t *ASTTool) Name() string { return "ast" } | ||
| 53 | func (t *ASTTool) Description() string { | 55 | func (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." | 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." | |||
| 55 | } | 64 | } | ||
| 56 | func (t *ASTTool) Schema() json.RawMessage { return json.RawMessage(astSchema) } | 65 | func (t *ASTTool) Schema() json.RawMessage { return json.RawMessage(astSchema) } | ||
| 57 | 66 | ||||
| @@ -79,10 +88,8 @@ func (t *ASTTool) Execute(ctx context.Context, raw json.RawMessage, _ func(strin | |||||
| 79 | if args.Rewrite == "" { | 88 | if args.Rewrite == "" { | ||
| 80 | return plan.result, nil | 89 | return plan.result, nil | ||
| 81 | } | 90 | } | ||
| 82 | - | for _, change := range plan.changes { | 91 | + | if err := t.applyRewrite(plan.changes); err != nil { |
| 83 | - | if err := os.WriteFile(change.path, change.content, change.mode); err != nil { | 92 | + | return core.ToolResult{}, err |
| 84 | - | return core.ToolResult{}, fmt.Errorf("ast: write %s: %w", change.path, err) | |||
| 85 | - | } | |||
| 86 | } | 93 | } | ||
| 87 | if t.LSPDiagnostics { | 94 | if t.LSPDiagnostics { | ||
| 88 | attachMutationDiagnostics(ctx, t.CWD, plan.result.Context.Mutates, t.LSP, &plan.result) | 95 | attachMutationDiagnostics(ctx, t.CWD, plan.result.Context.Mutates, t.LSP, &plan.result) | ||
| @@ -90,10 +97,90 @@ func (t *ASTTool) Execute(ctx context.Context, raw json.RawMessage, _ func(strin | |||||
| 90 | return plan.result, nil | 97 | return plan.result, nil | ||
| 91 | } | 98 | } | ||
| 92 | 99 | ||||
| 100 | + | func (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. | |||
| 129 | + | func (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 | + | ||||
| 93 | type astChange struct { | 152 | type astChange struct { | ||
| 94 | - | path string | 153 | + | path string |
| 95 | - | content []byte | 154 | + | original []byte |
| 96 | - | mode os.FileMode | 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. | |||
| 161 | + | func 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 | |||
| 97 | } | 184 | } | ||
| 98 | 185 | ||||
| 99 | type astPlan struct { | 186 | type astPlan struct { | ||
| @@ -124,8 +211,8 @@ func parseASTArgs(raw json.RawMessage) (astArgs, error) { | |||||
| 124 | } | 211 | } | ||
| 125 | args.Pattern = strings.TrimSpace(args.Pattern) | 212 | args.Pattern = strings.TrimSpace(args.Pattern) | ||
| 126 | args.Language = strings.TrimSpace(args.Language) | 213 | args.Language = strings.TrimSpace(args.Language) | ||
| 127 | - | if args.Pattern == "" || args.Language == "" || strings.TrimSpace(args.Path) == "" { | 214 | + | if args.Pattern == "" || strings.TrimSpace(args.Path) == "" { |
| 128 | - | return args, fmt.Errorf("ast: pattern, language, and path are required") | 215 | + | return args, fmt.Errorf("ast: pattern and path are required") |
| 129 | } | 216 | } | ||
| 130 | return args, nil | 217 | return args, nil | ||
| 131 | } | 218 | } | ||
| @@ -156,57 +243,69 @@ func (t *ASTTool) plan(ctx context.Context, args astArgs) (astPlan, error) { | |||||
| 156 | if err != nil { | 243 | if err != nil { | ||
| 157 | return astPlan{}, err | 244 | return astPlan{}, err | ||
| 158 | } | 245 | } | ||
| 159 | - | commandArgs := []string{"run", "--pattern=" + args.Pattern, "--lang=" + args.Language} | 246 | + | dir := root |
| 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() { | 247 | if !info.IsDir() { | ||
| 167 | - | cmd.Dir = filepath.Dir(root) | 248 | + | 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 | } | 249 | } | ||
| 178 | - | if capture.overflow { | 250 | + | commandArgs := []string{"run", "--pattern=" + args.Pattern} |
| 179 | - | return astPlan{}, fmt.Errorf("ast: output exceeds %d bytes; use a narrower path or pattern", capture.limit) | 251 | + | if args.Language != "" { |
| 252 | + | commandArgs = append(commandArgs, "--lang="+args.Language) | |||
| 180 | } | 253 | } | ||
| 181 | - | if runErr != nil { | 254 | + | if args.Rewrite != "" { |
| 182 | - | var exit *exec.ExitError | 255 | + | commandArgs = append(commandArgs, "--rewrite="+args.Rewrite) |
| 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 | } | 256 | } | ||
| 191 | - | matches, err := decodeASTMatches(stdout.Bytes()) | 257 | + | commandArgs = append(commandArgs, "--json=stream", root) |
| 258 | + | run, err := runASTGrep(ctx, binary, dir, commandArgs, args.Rewrite != "") | |||
| 192 | if err != nil { | 259 | if err != nil { | ||
| 193 | return astPlan{}, err | 260 | return astPlan{}, err | ||
| 194 | } | 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 | |||
| 195 | for index := range matches { | 276 | for index := range matches { | ||
| 196 | if !filepath.IsAbs(matches[index].File) { | 277 | if !filepath.IsAbs(matches[index].File) { | ||
| 197 | - | matches[index].File = filepath.Join(cmd.Dir, matches[index].File) | 278 | + | matches[index].File = filepath.Join(dir, matches[index].File) |
| 198 | } | 279 | } | ||
| 199 | matches[index].File = filepath.Clean(matches[index].File) | 280 | matches[index].File = filepath.Clean(matches[index].File) | ||
| 200 | } | 281 | } | ||
| 201 | if args.Rewrite == "" { | 282 | if args.Rewrite == "" { | ||
| 202 | text, paths := formatASTMatches(matches) | 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 | + | } | |||
| 203 | return astPlan{result: core.ToolResult{ | 290 | return astPlan{result: core.ToolResult{ | ||
| 204 | Content: []provider.Content{provider.TextBlock{Text: text}}, | 291 | Content: []provider.Content{provider.TextBlock{Text: text}}, | ||
| 205 | Context: provider.ToolContext{Discovers: paths}, | 292 | Context: provider.ToolContext{Discovers: paths}, | ||
| 206 | Details: map[string]any{"engine": "ast-grep", "path": root, "language": args.Language, "matches": len(matches)}, | 293 | Details: map[string]any{"engine": "ast-grep", "path": root, "language": args.Language, "matches": len(matches)}, | ||
| 207 | }}, nil | 294 | }}, nil | ||
| 208 | } | 295 | } | ||
| 209 | - | return t.planRewrite(root, args, matches) | 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. | |||
| 306 | + | func (t *ASTTool) Available() bool { | |||
| 307 | + | _, err := t.executable() | |||
| 308 | + | return err == nil | |||
| 210 | } | 309 | } | ||
| 211 | 310 | ||||
| 212 | func (t *ASTTool) executable() (string, error) { | 311 | func (t *ASTTool) executable() (string, error) { | ||
| @@ -214,14 +313,103 @@ func (t *ASTTool) executable() (string, error) { | |||||
| 214 | if lookup == nil { | 313 | if lookup == nil { | ||
| 215 | lookup = exec.LookPath | 314 | lookup = exec.LookPath | ||
| 216 | } | 315 | } | ||
| 217 | - | for _, name := range []string{"ast-grep", "sg"} { | 316 | + | if path, err := lookup("ast-grep"); err == nil && path != "" { |
| 218 | - | if path, err := lookup(name); 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") { | |||
| 219 | return path, nil | 324 | return path, nil | ||
| 220 | } | 325 | } | ||
| 221 | } | 326 | } | ||
| 222 | return "", fmt.Errorf("ast: ast-grep executable ('ast-grep' or 'sg') not found in PATH") | 327 | return "", fmt.Errorf("ast: ast-grep executable ('ast-grep' or 'sg') not found in PATH") | ||
| 223 | } | 328 | } | ||
| 224 | 329 | ||||
| 330 | + | const 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 | + | ||||
| 332 | + | type 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. | |||
| 340 | + | func 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 | + | ||||
| 380 | + | var 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. | |||
| 384 | + | func 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. | |||
| 397 | + | func 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 | + | ||||
| 225 | type limitedASTBuffer struct { | 413 | type limitedASTBuffer struct { | ||
| 226 | buffer *bytes.Buffer | 414 | buffer *bytes.Buffer | ||
| 227 | limit int | 415 | limit int | ||
| @@ -281,6 +469,17 @@ func formatASTMatches(matches []astMatch) (string, []string) { | |||||
| 281 | return boundASTText(strings.TrimSpace(output.String())), paths | 469 | return boundASTText(strings.TrimSpace(output.String())), paths | ||
| 282 | } | 470 | } | ||
| 283 | 471 | ||||
| 472 | + | func 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 | + | ||||
| 284 | func (t *ASTTool) planRewrite(root string, args astArgs, matches []astMatch) (astPlan, error) { | 483 | func (t *ASTTool) planRewrite(root string, args astArgs, matches []astMatch) (astPlan, error) { | ||
| 285 | byFile := make(map[string][]astMatch) | 484 | byFile := make(map[string][]astMatch) | ||
| 286 | for _, match := range matches { | 485 | for _, match := range matches { | ||
| @@ -342,20 +541,28 @@ func (t *ASTTool) planRewrite(root string, args astArgs, matches []astMatch) (as | |||||
| 342 | if bytes.Equal(original, updated) { | 541 | if bytes.Equal(original, updated) { | ||
| 343 | continue | 542 | continue | ||
| 344 | } | 543 | } | ||
| 345 | - | if strings.EqualFold(args.Language, "go") { | 544 | + | if strings.EqualFold(args.Language, "go") || strings.EqualFold(args.Language, "golang") || (args.Language == "" && strings.EqualFold(filepath.Ext(path), ".go")) { |
| 346 | - | formatted, formatErr := format.Source(updated) | 545 | + | crlf := bytes.Contains(original, []byte("\r\n")) |
| 347 | - | if formatErr == nil { | 546 | + | normalized := bytes.ReplaceAll(original, []byte("\r\n"), []byte("\n")) |
| 348 | - | if bytes.Contains(original, []byte("\r\n")) { | 547 | + | // Mixed endings and pre-existing gofmt differences must not cause |
| 349 | - | formatted = bytes.ReplaceAll(formatted, []byte("\n"), []byte("\r\n")) | 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 | + | } | |||
| 350 | } | 558 | } | ||
| 351 | - | updated = formatted | |||
| 352 | } | 559 | } | ||
| 353 | } | 560 | } | ||
| 354 | if bytes.Equal(original, updated) { | 561 | if bytes.Equal(original, updated) { | ||
| 355 | continue | 562 | continue | ||
| 356 | } | 563 | } | ||
| 357 | diffs.WriteString(unifiedDiff(path, string(original), string(updated))) | 564 | diffs.WriteString(unifiedDiff(path, string(original), string(updated))) | ||
| 358 | - | changes = append(changes, astChange{path: path, content: updated, mode: info.Mode().Perm()}) | 565 | + | changes = append(changes, astChange{path: path, original: original, content: updated, mode: info.Mode().Perm()}) |
| 359 | } | 566 | } | ||
| 360 | text := boundASTText(diffs.String()) | 567 | text := boundASTText(diffs.String()) | ||
| 361 | if len(changes) == 0 { | 568 | if len(changes) == 0 { | ||
M packages/agent/tools/ast_test.go +164 -7
| @@ -1,6 +1,7 @@ | |||
| 1 | 1 | package tools | |
| 2 | 2 | ||
| 3 | 3 | import ( | |
| 4 | + | "bytes" | |
| 4 | 5 | "context" | |
| 5 | 6 | "encoding/json" | |
| 6 | 7 | "os" | |
| @@ -27,15 +28,19 @@ func TestASTToolReportsCaptureLimitWithoutBrokenPipe(t *testing.T) { | |||
| 27 | 28 | t.Skip("test fixture uses a POSIX script") | |
| 28 | 29 | } | |
| 29 | 30 | root := t.TempDir() | |
| 30 | - | script := filepath.Join(root, "sg") | |
| 31 | + | script := filepath.Join(root, "ast-grep") | |
| 31 | 32 | body := "#!/bin/sh\nhead -c 4194305 /dev/zero | tr '\\0' x\n" | |
| 32 | 33 | if err := os.WriteFile(script, []byte(body), 0o755); err != nil { | |
| 33 | 34 | t.Fatal(err) | |
| 34 | 35 | } | |
| 35 | 36 | tool := &ASTTool{CWD: root, LookPath: func(string) (string, error) { return script, nil }} | |
| 36 | - | _, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":"."}`), nil) | |
| 37 | + | result, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":"."}`), nil) | |
| 38 | + | if err != nil || !strings.Contains(result.Content[0].(provider.TextBlock).Text, "search output truncated") { | |
| 39 | + | t.Fatalf("search result = %+v, error = %v", result, err) | |
| 40 | + | } | |
| 41 | + | _, err = tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":".","rewrite":"x"}`), nil) | |
| 37 | 42 | if err == nil || !strings.Contains(err.Error(), "output exceeds") || strings.Contains(strings.ToLower(err.Error()), "broken pipe") { | |
| 38 | - | t.Fatalf("Execute error = %v", err) | |
| 43 | + | t.Fatalf("rewrite error = %v", err) | |
| 39 | 44 | } | |
| 40 | 45 | } | |
| 41 | 46 | ||
| @@ -44,7 +49,7 @@ func TestASTToolExecutesWithoutShellAndTracksDiscoveredFiles(t *testing.T) { | |||
| 44 | 49 | t.Skip("test fixture uses a POSIX script") | |
| 45 | 50 | } | |
| 46 | 51 | root := t.TempDir() | |
| 47 | - | script := filepath.Join(root, "sg") | |
| 52 | + | script := filepath.Join(root, "ast-grep") | |
| 48 | 53 | body := "#!/bin/sh\nprintf '%s\\n' '{\"file\":\"main.go\",\"range\":{}}'\n" | |
| 49 | 54 | if err := os.WriteFile(script, []byte(body), 0o755); err != nil { | |
| 50 | 55 | t.Fatal(err) | |
| @@ -107,7 +112,7 @@ func TestASTToolRewritePreviewAndExecute(t *testing.T) { | |||
| 107 | 112 | if !strings.Contains(string(got), "fmt.Println") || len(result.Context.Mutates) != 1 { | |
| 108 | 113 | t.Fatalf("result = %q, context = %+v", got, result.Context) | |
| 109 | 114 | } | |
| 110 | - | if info, _ := os.Stat(path); info.Mode().Perm() != 0o640 { | |
| 115 | + | if info, _ := os.Stat(path); runtime.GOOS != "windows" && info.Mode().Perm() != 0o640 { | |
| 111 | 116 | t.Fatalf("mode = %v", info.Mode().Perm()) | |
| 112 | 117 | } | |
| 113 | 118 | } | |
| @@ -115,11 +120,12 @@ func TestASTToolRewritePreviewAndExecute(t *testing.T) { | |||
| 115 | 120 | func TestASTRewriteFormatsGoSource(t *testing.T) { | |
| 116 | 121 | root := t.TempDir() | |
| 117 | 122 | path := filepath.Join(root, "main.go") | |
| 118 | - | original := []byte("package main\nfunc add(left, right int) int { return left + right\n }\n") | |
| 123 | + | original := []byte("package main\n\nfunc add(left, right int) int { return left + right }\n") | |
| 119 | 124 | if err := os.WriteFile(path, original, 0o644); err != nil { | |
| 120 | 125 | t.Fatal(err) | |
| 121 | 126 | } | |
| 122 | - | match := astTestMatch(path, "return left + right\n ", "\nreturn left + right\n\n", 45, 66) | |
| 127 | + | start := bytes.Index(original, []byte("return left + right")) | |
| 128 | + | match := astTestMatch(path, "return left + right", "return left + right\n\n", start, start+len("return left + right")) | |
| 123 | 129 | plan, err := (&ASTTool{CWD: root}).planRewrite(root, astArgs{Language: "go", Rewrite: "{$$$BODY}"}, []astMatch{match}) | |
| 124 | 130 | if err != nil { | |
| 125 | 131 | t.Fatalf("planRewrite: %v", err) | |
| @@ -163,3 +169,154 @@ func TestBoundASTTextPreservesUTF8(t *testing.T) { | |||
| 163 | 169 | t.Fatalf("invalid bounded output") | |
| 164 | 170 | } | |
| 165 | 171 | } | |
| 172 | + | ||
| 173 | + | func TestASTFallbackRejectsOtherSG(t *testing.T) { | |
| 174 | + | if runtime.GOOS == "windows" { | |
| 175 | + | t.Skip("POSIX script") | |
| 176 | + | } | |
| 177 | + | root := t.TempDir() | |
| 178 | + | path := filepath.Join(root, "sg") | |
| 179 | + | if err := os.WriteFile(path, []byte("#!/bin/sh\necho 'shadow-utils 1.0'\n"), 0o755); err != nil { | |
| 180 | + | t.Fatal(err) | |
| 181 | + | } | |
| 182 | + | tool := &ASTTool{LookPath: func(name string) (string, error) { | |
| 183 | + | if name == "sg" { | |
| 184 | + | return path, nil | |
| 185 | + | } | |
| 186 | + | return "", os.ErrNotExist | |
| 187 | + | }} | |
| 188 | + | if _, err := tool.executable(); err == nil { | |
| 189 | + | t.Fatal("accepted non ast-grep sg") | |
| 190 | + | } | |
| 191 | + | if err := os.WriteFile(path, []byte("#!/bin/sh\necho 'ast-grep 0.44.1'\n"), 0o755); err != nil { | |
| 192 | + | t.Fatal(err) | |
| 193 | + | } | |
| 194 | + | if got, err := tool.executable(); err != nil || got != path { | |
| 195 | + | t.Fatalf("executable = %q, %v", got, err) | |
| 196 | + | } | |
| 197 | + | } | |
| 198 | + | ||
| 199 | + | func TestASTRewriteFormattingDependsOnOriginal(t *testing.T) { | |
| 200 | + | for _, tc := range []struct{ name, original, replacement, want string }{ | |
| 201 | + | {"unclean", "package main\nfunc main(){println(1)}\n", "println(2)", "package main\nfunc main(){println(2)}\n"}, | |
| 202 | + | {"crlf", "package main\r\n\r\nfunc main() { println(1) }\r\n", "println( 2 )", "package main\r\n\r\nfunc main() { println(2) }\r\n"}, | |
| 203 | + | {"mixed", "package main\r\n\nfunc main() { println(1) }\n", "println( 2 )", "package main\r\n\nfunc main() { println( 2 ) }\n"}, | |
| 204 | + | } { | |
| 205 | + | t.Run(tc.name, func(t *testing.T) { | |
| 206 | + | root := t.TempDir() | |
| 207 | + | path := filepath.Join(root, "main.go") | |
| 208 | + | if err := os.WriteFile(path, []byte(tc.original), 0o644); err != nil { | |
| 209 | + | t.Fatal(err) | |
| 210 | + | } | |
| 211 | + | start := strings.Index(tc.original, "println(1)") | |
| 212 | + | match := astTestMatch(path, "println(1)", tc.replacement, start, start+len("println(1)")) | |
| 213 | + | plan, err := (&ASTTool{CWD: root}).planRewrite(root, astArgs{Language: "go", Rewrite: "x"}, []astMatch{match}) | |
| 214 | + | if err != nil { | |
| 215 | + | t.Fatal(err) | |
| 216 | + | } | |
| 217 | + | if got := string(plan.changes[0].content); got != tc.want { | |
| 218 | + | t.Fatalf("got %q, want %q", got, tc.want) | |
| 219 | + | } | |
| 220 | + | }) | |
| 221 | + | } | |
| 222 | + | } | |
| 223 | + | ||
| 224 | + | func TestASTScopeAndReadOnly(t *testing.T) { | |
| 225 | + | root := t.TempDir() | |
| 226 | + | outside := t.TempDir() | |
| 227 | + | path := filepath.Join(outside, "main.go") | |
| 228 | + | if err := os.WriteFile(path, []byte("package main\n"), 0o644); err != nil { | |
| 229 | + | t.Fatal(err) | |
| 230 | + | } | |
| 231 | + | sandbox := NewSandbox(root) | |
| 232 | + | sandbox.Lock() | |
| 233 | + | tool := &ASTTool{CWD: root, Sandbox: sandbox} | |
| 234 | + | // Marshal the path: raw Windows backslashes are invalid JSON escapes. | |
| 235 | + | args, err := json.Marshal(map[string]string{"pattern": "$A", "language": "go", "path": path}) | |
| 236 | + | if err != nil { | |
| 237 | + | t.Fatal(err) | |
| 238 | + | } | |
| 239 | + | if _, err := tool.Execute(context.Background(), args, nil); err == nil || !strings.Contains(err.Error(), "outside sandbox") { | |
| 240 | + | t.Fatalf("jail: %v", err) | |
| 241 | + | } | |
| 242 | + | tool.ReadOnly = true | |
| 243 | + | if _, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":".","rewrite":"x"}`), nil); err == nil || !strings.Contains(err.Error(), "read-only") { | |
| 244 | + | t.Fatalf("read-only: %v", err) | |
| 245 | + | } | |
| 246 | + | tool.ReadOnly = false | |
| 247 | + | match := astTestMatch(path, "package", "module", 0, 7) | |
| 248 | + | if _, err := tool.planRewrite(root, astArgs{Rewrite: "module"}, []astMatch{match}); err == nil || !strings.Contains(err.Error(), "outside sandbox") { | |
| 249 | + | t.Fatalf("write scope: %v", err) | |
| 250 | + | } | |
| 251 | + | permissions := &PermissionSet{} | |
| 252 | + | permissions.FS.Read = []string{root, outside} | |
| 253 | + | permissions.FS.Write = []string{root} | |
| 254 | + | unlocked := NewSandbox(root) | |
| 255 | + | unlocked.SetPermissions(permissions) | |
| 256 | + | if _, err := (&ASTTool{CWD: root, Sandbox: unlocked}).planRewrite(root, astArgs{Rewrite: "module"}, []astMatch{match}); err == nil || !strings.Contains(err.Error(), "outside declared scopes") { | |
| 257 | + | t.Fatalf("permission: %v", err) | |
| 258 | + | } | |
| 259 | + | } | |
| 260 | + | ||
| 261 | + | func TestASTWarningsAndExitError(t *testing.T) { | |
| 262 | + | if runtime.GOOS == "windows" { | |
| 263 | + | t.Skip("POSIX script") | |
| 264 | + | } | |
| 265 | + | root := t.TempDir() | |
| 266 | + | script := filepath.Join(root, "ast-grep") | |
| 267 | + | tool := &ASTTool{CWD: root, LookPath: func(string) (string, error) { return script, nil }} | |
| 268 | + | if err := os.WriteFile(script, []byte("#!/bin/sh\necho 'Pattern contains an ERROR node' >&2\n"), 0o755); err != nil { | |
| 269 | + | t.Fatal(err) | |
| 270 | + | } | |
| 271 | + | for _, rewrite := range []string{"", `,"rewrite":"x"`} { | |
| 272 | + | result, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":"."`+rewrite+`}`), nil) | |
| 273 | + | if err != nil || !strings.Contains(result.Content[0].(provider.TextBlock).Text, "No structural matches.\nPattern contains an ERROR node") { | |
| 274 | + | t.Fatalf("result = %+v, %v", result, err) | |
| 275 | + | } | |
| 276 | + | } | |
| 277 | + | if err := os.WriteFile(script, []byte("#!/bin/sh\necho 'bad pattern' >&2\nexit 2\n"), 0o755); err != nil { | |
| 278 | + | t.Fatal(err) | |
| 279 | + | } | |
| 280 | + | if _, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":"."}`), nil); err == nil || !strings.Contains(err.Error(), "bad pattern") { | |
| 281 | + | t.Fatalf("exit error: %v", err) | |
| 282 | + | } | |
| 283 | + | } | |
| 284 | + | ||
| 285 | + | func TestASTApplyRewriteValidatesAllBeforeWriting(t *testing.T) { | |
| 286 | + | root := t.TempDir() | |
| 287 | + | first := filepath.Join(root, "a.go") | |
| 288 | + | second := filepath.Join(root, "b.go") | |
| 289 | + | for _, path := range []string{first, second} { | |
| 290 | + | if err := os.WriteFile(path, []byte("old"), 0o640); err != nil { | |
| 291 | + | t.Fatal(err) | |
| 292 | + | } | |
| 293 | + | } | |
| 294 | + | changes := []astChange{ | |
| 295 | + | {path: first, original: []byte("old"), content: []byte("new"), mode: 0o640}, | |
| 296 | + | {path: second, original: []byte("old"), content: []byte("new"), mode: 0o640}, | |
| 297 | + | } | |
| 298 | + | if err := os.WriteFile(second, []byte("changed"), 0o640); err != nil { | |
| 299 | + | t.Fatal(err) | |
| 300 | + | } | |
| 301 | + | if err := (&ASTTool{}).applyRewrite(changes); err == nil || !strings.Contains(err.Error(), "changed while applying") { | |
| 302 | + | t.Fatalf("error = %v", err) | |
| 303 | + | } | |
| 304 | + | if got, _ := os.ReadFile(first); string(got) != "old" { | |
| 305 | + | t.Fatalf("first file changed: %q", got) | |
| 306 | + | } | |
| 307 | + | if err := os.WriteFile(second, []byte("old"), 0o640); err != nil { | |
| 308 | + | t.Fatal(err) | |
| 309 | + | } | |
| 310 | + | if err := (&ASTTool{}).applyRewrite(changes); err != nil { | |
| 311 | + | t.Fatal(err) | |
| 312 | + | } | |
| 313 | + | if got, _ := os.ReadFile(first); string(got) != "new" { | |
| 314 | + | t.Fatalf("first file = %q", got) | |
| 315 | + | } | |
| 316 | + | // Windows has no Unix permission bits to preserve. | |
| 317 | + | if runtime.GOOS != "windows" { | |
| 318 | + | if info, err := os.Stat(first); err != nil || info.Mode().Perm() != 0o640 { | |
| 319 | + | t.Fatalf("mode = %v, %v", info, err) | |
| 320 | + | } | |
| 321 | + | } | |
| 322 | + | } | |
| @@ -1,6 +1,7 @@ | |||||
| 1 | package tools | 1 | package tools | ||
| 2 | 2 | ||||
| 3 | import ( | 3 | import ( | ||
| 4 | + | "bytes" | |||
| 4 | "context" | 5 | "context" | ||
| 5 | "encoding/json" | 6 | "encoding/json" | ||
| 6 | "os" | 7 | "os" | ||
| @@ -27,15 +28,19 @@ func TestASTToolReportsCaptureLimitWithoutBrokenPipe(t *testing.T) { | |||||
| 27 | t.Skip("test fixture uses a POSIX script") | 28 | t.Skip("test fixture uses a POSIX script") | ||
| 28 | } | 29 | } | ||
| 29 | root := t.TempDir() | 30 | root := t.TempDir() | ||
| 30 | - | script := filepath.Join(root, "sg") | 31 | + | script := filepath.Join(root, "ast-grep") |
| 31 | body := "#!/bin/sh\nhead -c 4194305 /dev/zero | tr '\\0' x\n" | 32 | body := "#!/bin/sh\nhead -c 4194305 /dev/zero | tr '\\0' x\n" | ||
| 32 | if err := os.WriteFile(script, []byte(body), 0o755); err != nil { | 33 | if err := os.WriteFile(script, []byte(body), 0o755); err != nil { | ||
| 33 | t.Fatal(err) | 34 | t.Fatal(err) | ||
| 34 | } | 35 | } | ||
| 35 | tool := &ASTTool{CWD: root, LookPath: func(string) (string, error) { return script, nil }} | 36 | tool := &ASTTool{CWD: root, LookPath: func(string) (string, error) { return script, nil }} | ||
| 36 | - | _, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":"."}`), nil) | 37 | + | result, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":"."}`), nil) |
| 38 | + | if err != nil || !strings.Contains(result.Content[0].(provider.TextBlock).Text, "search output truncated") { | |||
| 39 | + | t.Fatalf("search result = %+v, error = %v", result, err) | |||
| 40 | + | } | |||
| 41 | + | _, err = tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":".","rewrite":"x"}`), nil) | |||
| 37 | if err == nil || !strings.Contains(err.Error(), "output exceeds") || strings.Contains(strings.ToLower(err.Error()), "broken pipe") { | 42 | if err == nil || !strings.Contains(err.Error(), "output exceeds") || strings.Contains(strings.ToLower(err.Error()), "broken pipe") { | ||
| 38 | - | t.Fatalf("Execute error = %v", err) | 43 | + | t.Fatalf("rewrite error = %v", err) |
| 39 | } | 44 | } | ||
| 40 | } | 45 | } | ||
| 41 | 46 | ||||
| @@ -44,7 +49,7 @@ func TestASTToolExecutesWithoutShellAndTracksDiscoveredFiles(t *testing.T) { | |||||
| 44 | t.Skip("test fixture uses a POSIX script") | 49 | t.Skip("test fixture uses a POSIX script") | ||
| 45 | } | 50 | } | ||
| 46 | root := t.TempDir() | 51 | root := t.TempDir() | ||
| 47 | - | script := filepath.Join(root, "sg") | 52 | + | script := filepath.Join(root, "ast-grep") |
| 48 | body := "#!/bin/sh\nprintf '%s\\n' '{\"file\":\"main.go\",\"range\":{}}'\n" | 53 | body := "#!/bin/sh\nprintf '%s\\n' '{\"file\":\"main.go\",\"range\":{}}'\n" | ||
| 49 | if err := os.WriteFile(script, []byte(body), 0o755); err != nil { | 54 | if err := os.WriteFile(script, []byte(body), 0o755); err != nil { | ||
| 50 | t.Fatal(err) | 55 | t.Fatal(err) | ||
| @@ -107,7 +112,7 @@ func TestASTToolRewritePreviewAndExecute(t *testing.T) { | |||||
| 107 | if !strings.Contains(string(got), "fmt.Println") || len(result.Context.Mutates) != 1 { | 112 | if !strings.Contains(string(got), "fmt.Println") || len(result.Context.Mutates) != 1 { | ||
| 108 | t.Fatalf("result = %q, context = %+v", got, result.Context) | 113 | t.Fatalf("result = %q, context = %+v", got, result.Context) | ||
| 109 | } | 114 | } | ||
| 110 | - | if info, _ := os.Stat(path); info.Mode().Perm() != 0o640 { | 115 | + | if info, _ := os.Stat(path); runtime.GOOS != "windows" && info.Mode().Perm() != 0o640 { |
| 111 | t.Fatalf("mode = %v", info.Mode().Perm()) | 116 | t.Fatalf("mode = %v", info.Mode().Perm()) | ||
| 112 | } | 117 | } | ||
| 113 | } | 118 | } | ||
| @@ -115,11 +120,12 @@ func TestASTToolRewritePreviewAndExecute(t *testing.T) { | |||||
| 115 | func TestASTRewriteFormatsGoSource(t *testing.T) { | 120 | func TestASTRewriteFormatsGoSource(t *testing.T) { | ||
| 116 | root := t.TempDir() | 121 | root := t.TempDir() | ||
| 117 | path := filepath.Join(root, "main.go") | 122 | path := filepath.Join(root, "main.go") | ||
| 118 | - | original := []byte("package main\nfunc add(left, right int) int { return left + right\n }\n") | 123 | + | original := []byte("package main\n\nfunc add(left, right int) int { return left + right }\n") |
| 119 | if err := os.WriteFile(path, original, 0o644); err != nil { | 124 | if err := os.WriteFile(path, original, 0o644); err != nil { | ||
| 120 | t.Fatal(err) | 125 | t.Fatal(err) | ||
| 121 | } | 126 | } | ||
| 122 | - | match := astTestMatch(path, "return left + right\n ", "\nreturn left + right\n\n", 45, 66) | 127 | + | start := bytes.Index(original, []byte("return left + right")) |
| 128 | + | match := astTestMatch(path, "return left + right", "return left + right\n\n", start, start+len("return left + right")) | |||
| 123 | plan, err := (&ASTTool{CWD: root}).planRewrite(root, astArgs{Language: "go", Rewrite: "{$$$BODY}"}, []astMatch{match}) | 129 | plan, err := (&ASTTool{CWD: root}).planRewrite(root, astArgs{Language: "go", Rewrite: "{$$$BODY}"}, []astMatch{match}) | ||
| 124 | if err != nil { | 130 | if err != nil { | ||
| 125 | t.Fatalf("planRewrite: %v", err) | 131 | t.Fatalf("planRewrite: %v", err) | ||
| @@ -163,3 +169,154 @@ func TestBoundASTTextPreservesUTF8(t *testing.T) { | |||||
| 163 | t.Fatalf("invalid bounded output") | 169 | t.Fatalf("invalid bounded output") | ||
| 164 | } | 170 | } | ||
| 165 | } | 171 | } | ||
| 172 | + | ||||
| 173 | + | func TestASTFallbackRejectsOtherSG(t *testing.T) { | |||
| 174 | + | if runtime.GOOS == "windows" { | |||
| 175 | + | t.Skip("POSIX script") | |||
| 176 | + | } | |||
| 177 | + | root := t.TempDir() | |||
| 178 | + | path := filepath.Join(root, "sg") | |||
| 179 | + | if err := os.WriteFile(path, []byte("#!/bin/sh\necho 'shadow-utils 1.0'\n"), 0o755); err != nil { | |||
| 180 | + | t.Fatal(err) | |||
| 181 | + | } | |||
| 182 | + | tool := &ASTTool{LookPath: func(name string) (string, error) { | |||
| 183 | + | if name == "sg" { | |||
| 184 | + | return path, nil | |||
| 185 | + | } | |||
| 186 | + | return "", os.ErrNotExist | |||
| 187 | + | }} | |||
| 188 | + | if _, err := tool.executable(); err == nil { | |||
| 189 | + | t.Fatal("accepted non ast-grep sg") | |||
| 190 | + | } | |||
| 191 | + | if err := os.WriteFile(path, []byte("#!/bin/sh\necho 'ast-grep 0.44.1'\n"), 0o755); err != nil { | |||
| 192 | + | t.Fatal(err) | |||
| 193 | + | } | |||
| 194 | + | if got, err := tool.executable(); err != nil || got != path { | |||
| 195 | + | t.Fatalf("executable = %q, %v", got, err) | |||
| 196 | + | } | |||
| 197 | + | } | |||
| 198 | + | ||||
| 199 | + | func TestASTRewriteFormattingDependsOnOriginal(t *testing.T) { | |||
| 200 | + | for _, tc := range []struct{ name, original, replacement, want string }{ | |||
| 201 | + | {"unclean", "package main\nfunc main(){println(1)}\n", "println(2)", "package main\nfunc main(){println(2)}\n"}, | |||
| 202 | + | {"crlf", "package main\r\n\r\nfunc main() { println(1) }\r\n", "println( 2 )", "package main\r\n\r\nfunc main() { println(2) }\r\n"}, | |||
| 203 | + | {"mixed", "package main\r\n\nfunc main() { println(1) }\n", "println( 2 )", "package main\r\n\nfunc main() { println( 2 ) }\n"}, | |||
| 204 | + | } { | |||
| 205 | + | t.Run(tc.name, func(t *testing.T) { | |||
| 206 | + | root := t.TempDir() | |||
| 207 | + | path := filepath.Join(root, "main.go") | |||
| 208 | + | if err := os.WriteFile(path, []byte(tc.original), 0o644); err != nil { | |||
| 209 | + | t.Fatal(err) | |||
| 210 | + | } | |||
| 211 | + | start := strings.Index(tc.original, "println(1)") | |||
| 212 | + | match := astTestMatch(path, "println(1)", tc.replacement, start, start+len("println(1)")) | |||
| 213 | + | plan, err := (&ASTTool{CWD: root}).planRewrite(root, astArgs{Language: "go", Rewrite: "x"}, []astMatch{match}) | |||
| 214 | + | if err != nil { | |||
| 215 | + | t.Fatal(err) | |||
| 216 | + | } | |||
| 217 | + | if got := string(plan.changes[0].content); got != tc.want { | |||
| 218 | + | t.Fatalf("got %q, want %q", got, tc.want) | |||
| 219 | + | } | |||
| 220 | + | }) | |||
| 221 | + | } | |||
| 222 | + | } | |||
| 223 | + | ||||
| 224 | + | func TestASTScopeAndReadOnly(t *testing.T) { | |||
| 225 | + | root := t.TempDir() | |||
| 226 | + | outside := t.TempDir() | |||
| 227 | + | path := filepath.Join(outside, "main.go") | |||
| 228 | + | if err := os.WriteFile(path, []byte("package main\n"), 0o644); err != nil { | |||
| 229 | + | t.Fatal(err) | |||
| 230 | + | } | |||
| 231 | + | sandbox := NewSandbox(root) | |||
| 232 | + | sandbox.Lock() | |||
| 233 | + | tool := &ASTTool{CWD: root, Sandbox: sandbox} | |||
| 234 | + | // Marshal the path: raw Windows backslashes are invalid JSON escapes. | |||
| 235 | + | args, err := json.Marshal(map[string]string{"pattern": "$A", "language": "go", "path": path}) | |||
| 236 | + | if err != nil { | |||
| 237 | + | t.Fatal(err) | |||
| 238 | + | } | |||
| 239 | + | if _, err := tool.Execute(context.Background(), args, nil); err == nil || !strings.Contains(err.Error(), "outside sandbox") { | |||
| 240 | + | t.Fatalf("jail: %v", err) | |||
| 241 | + | } | |||
| 242 | + | tool.ReadOnly = true | |||
| 243 | + | if _, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":".","rewrite":"x"}`), nil); err == nil || !strings.Contains(err.Error(), "read-only") { | |||
| 244 | + | t.Fatalf("read-only: %v", err) | |||
| 245 | + | } | |||
| 246 | + | tool.ReadOnly = false | |||
| 247 | + | match := astTestMatch(path, "package", "module", 0, 7) | |||
| 248 | + | if _, err := tool.planRewrite(root, astArgs{Rewrite: "module"}, []astMatch{match}); err == nil || !strings.Contains(err.Error(), "outside sandbox") { | |||
| 249 | + | t.Fatalf("write scope: %v", err) | |||
| 250 | + | } | |||
| 251 | + | permissions := &PermissionSet{} | |||
| 252 | + | permissions.FS.Read = []string{root, outside} | |||
| 253 | + | permissions.FS.Write = []string{root} | |||
| 254 | + | unlocked := NewSandbox(root) | |||
| 255 | + | unlocked.SetPermissions(permissions) | |||
| 256 | + | if _, err := (&ASTTool{CWD: root, Sandbox: unlocked}).planRewrite(root, astArgs{Rewrite: "module"}, []astMatch{match}); err == nil || !strings.Contains(err.Error(), "outside declared scopes") { | |||
| 257 | + | t.Fatalf("permission: %v", err) | |||
| 258 | + | } | |||
| 259 | + | } | |||
| 260 | + | ||||
| 261 | + | func TestASTWarningsAndExitError(t *testing.T) { | |||
| 262 | + | if runtime.GOOS == "windows" { | |||
| 263 | + | t.Skip("POSIX script") | |||
| 264 | + | } | |||
| 265 | + | root := t.TempDir() | |||
| 266 | + | script := filepath.Join(root, "ast-grep") | |||
| 267 | + | tool := &ASTTool{CWD: root, LookPath: func(string) (string, error) { return script, nil }} | |||
| 268 | + | if err := os.WriteFile(script, []byte("#!/bin/sh\necho 'Pattern contains an ERROR node' >&2\n"), 0o755); err != nil { | |||
| 269 | + | t.Fatal(err) | |||
| 270 | + | } | |||
| 271 | + | for _, rewrite := range []string{"", `,"rewrite":"x"`} { | |||
| 272 | + | result, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":"."`+rewrite+`}`), nil) | |||
| 273 | + | if err != nil || !strings.Contains(result.Content[0].(provider.TextBlock).Text, "No structural matches.\nPattern contains an ERROR node") { | |||
| 274 | + | t.Fatalf("result = %+v, %v", result, err) | |||
| 275 | + | } | |||
| 276 | + | } | |||
| 277 | + | if err := os.WriteFile(script, []byte("#!/bin/sh\necho 'bad pattern' >&2\nexit 2\n"), 0o755); err != nil { | |||
| 278 | + | t.Fatal(err) | |||
| 279 | + | } | |||
| 280 | + | if _, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"$A","language":"go","path":"."}`), nil); err == nil || !strings.Contains(err.Error(), "bad pattern") { | |||
| 281 | + | t.Fatalf("exit error: %v", err) | |||
| 282 | + | } | |||
| 283 | + | } | |||
| 284 | + | ||||
| 285 | + | func TestASTApplyRewriteValidatesAllBeforeWriting(t *testing.T) { | |||
| 286 | + | root := t.TempDir() | |||
| 287 | + | first := filepath.Join(root, "a.go") | |||
| 288 | + | second := filepath.Join(root, "b.go") | |||
| 289 | + | for _, path := range []string{first, second} { | |||
| 290 | + | if err := os.WriteFile(path, []byte("old"), 0o640); err != nil { | |||
| 291 | + | t.Fatal(err) | |||
| 292 | + | } | |||
| 293 | + | } | |||
| 294 | + | changes := []astChange{ | |||
| 295 | + | {path: first, original: []byte("old"), content: []byte("new"), mode: 0o640}, | |||
| 296 | + | {path: second, original: []byte("old"), content: []byte("new"), mode: 0o640}, | |||
| 297 | + | } | |||
| 298 | + | if err := os.WriteFile(second, []byte("changed"), 0o640); err != nil { | |||
| 299 | + | t.Fatal(err) | |||
| 300 | + | } | |||
| 301 | + | if err := (&ASTTool{}).applyRewrite(changes); err == nil || !strings.Contains(err.Error(), "changed while applying") { | |||
| 302 | + | t.Fatalf("error = %v", err) | |||
| 303 | + | } | |||
| 304 | + | if got, _ := os.ReadFile(first); string(got) != "old" { | |||
| 305 | + | t.Fatalf("first file changed: %q", got) | |||
| 306 | + | } | |||
| 307 | + | if err := os.WriteFile(second, []byte("old"), 0o640); err != nil { | |||
| 308 | + | t.Fatal(err) | |||
| 309 | + | } | |||
| 310 | + | if err := (&ASTTool{}).applyRewrite(changes); err != nil { | |||
| 311 | + | t.Fatal(err) | |||
| 312 | + | } | |||
| 313 | + | if got, _ := os.ReadFile(first); string(got) != "new" { | |||
| 314 | + | t.Fatalf("first file = %q", got) | |||
| 315 | + | } | |||
| 316 | + | // Windows has no Unix permission bits to preserve. | |||
| 317 | + | if runtime.GOOS != "windows" { | |||
| 318 | + | if info, err := os.Stat(first); err != nil || info.Mode().Perm() != 0o640 { | |||
| 319 | + | t.Fatalf("mode = %v, %v", info, err) | |||
| 320 | + | } | |||
| 321 | + | } | |||
| 322 | + | } | |||
A packages/agent/tools/ast_usability_test.go +186 -0
| @@ -0,0 +1,186 @@ | |||
| 1 | + | package tools | |
| 2 | + | ||
| 3 | + | import ( | |
| 4 | + | "context" | |
| 5 | + | "encoding/json" | |
| 6 | + | "os" | |
| 7 | + | "os/exec" | |
| 8 | + | "path/filepath" | |
| 9 | + | "runtime" | |
| 10 | + | "strings" | |
| 11 | + | "testing" | |
| 12 | + | ||
| 13 | + | "github.com/bnema/zut/packages/provider" | |
| 14 | + | ) | |
| 15 | + | ||
| 16 | + | func realASTTool(t *testing.T, root string) *ASTTool { | |
| 17 | + | t.Helper() | |
| 18 | + | binary, err := exec.LookPath("ast-grep") | |
| 19 | + | if err != nil { | |
| 20 | + | t.Skip("ast-grep not installed") | |
| 21 | + | } | |
| 22 | + | return &ASTTool{CWD: root, LookPath: func(string) (string, error) { return binary, nil }} | |
| 23 | + | } | |
| 24 | + | ||
| 25 | + | func astText(t *testing.T, tool *ASTTool, raw string) string { | |
| 26 | + | t.Helper() | |
| 27 | + | result, err := tool.Execute(context.Background(), json.RawMessage(raw), nil) | |
| 28 | + | if err != nil { | |
| 29 | + | t.Fatalf("Execute(%s): %v", raw, err) | |
| 30 | + | } | |
| 31 | + | return result.Content[0].(provider.TextBlock).Text | |
| 32 | + | } | |
| 33 | + | ||
| 34 | + | func TestASTLanguageIsOptional(t *testing.T) { | |
| 35 | + | root := t.TempDir() | |
| 36 | + | if err := os.WriteFile(filepath.Join(root, "a.py"), []byte("print(1)\n"), 0o644); err != nil { | |
| 37 | + | t.Fatal(err) | |
| 38 | + | } | |
| 39 | + | if err := os.WriteFile(filepath.Join(root, "b.ts"), []byte("console.log(2)\n"), 0o644); err != nil { | |
| 40 | + | t.Fatal(err) | |
| 41 | + | } | |
| 42 | + | tool := realASTTool(t, root) | |
| 43 | + | if got := astText(t, tool, `{"pattern":"print($A)","path":"."}`); !strings.Contains(got, "a.py:1") { | |
| 44 | + | t.Fatalf("python search = %q", got) | |
| 45 | + | } | |
| 46 | + | if got := astText(t, tool, `{"pattern":"console.log($$$)","path":"."}`); !strings.Contains(got, "b.ts:1") { | |
| 47 | + | t.Fatalf("typescript search = %q", got) | |
| 48 | + | } | |
| 49 | + | } | |
| 50 | + | ||
| 51 | + | // tree-sitter-go parses a bare `pkg.Func(args)` pattern as a type | |
| 52 | + | // conversion, so without the retry these calls never match. | |
| 53 | + | func TestASTGoQualifiedCallMatchesCalls(t *testing.T) { | |
| 54 | + | root := t.TempDir() | |
| 55 | + | path := filepath.Join(root, "main.go") | |
| 56 | + | source := "package main\n\nimport \"fmt\"\n\nfunc main() {\n\t_ = fmt.Errorf(\"a\")\n\t_ = fmt.Errorf(\"b %d\",\n\t\t1)\n}\n" | |
| 57 | + | if err := os.WriteFile(path, []byte(source), 0o644); err != nil { | |
| 58 | + | t.Fatal(err) | |
| 59 | + | } | |
| 60 | + | tool := realASTTool(t, root) | |
| 61 | + | for _, raw := range []string{ | |
| 62 | + | `{"pattern":"fmt.Errorf($$$ARGS)","language":"go","path":"."}`, | |
| 63 | + | `{"pattern":"fmt.Errorf($$$ARGS)","path":"."}`, | |
| 64 | + | `{"pattern":"fmt.Errorf($$$ARGS)","path":"main.go"}`, | |
| 65 | + | } { | |
| 66 | + | got := astText(t, tool, raw) | |
| 67 | + | if !strings.Contains(got, "main.go:6") || !strings.Contains(got, "main.go:7") { | |
| 68 | + | t.Fatalf("search %s = %q", raw, got) | |
| 69 | + | } | |
| 70 | + | } | |
| 71 | + | astText(t, tool, `{"pattern":"fmt.Errorf($$$ARGS)","path":".","rewrite":"errors.New($$$ARGS)"}`) | |
| 72 | + | got, _ := os.ReadFile(path) | |
| 73 | + | if strings.Contains(string(got), "fmt.Errorf") || strings.Count(string(got), "errors.New(") != 2 { | |
| 74 | + | t.Fatalf("rewrite = %q", got) | |
| 75 | + | } | |
| 76 | + | } | |
| 77 | + | ||
| 78 | + | func TestASTRewriteWritesThroughSymlink(t *testing.T) { | |
| 79 | + | root := t.TempDir() | |
| 80 | + | target := filepath.Join(root, "real.go") | |
| 81 | + | link := filepath.Join(root, "link.go") | |
| 82 | + | if err := os.WriteFile(target, []byte("old"), 0o644); err != nil { | |
| 83 | + | t.Fatal(err) | |
| 84 | + | } | |
| 85 | + | if err := os.Symlink(target, link); err != nil { | |
| 86 | + | t.Skip("symlinks unavailable:", err) | |
| 87 | + | } | |
| 88 | + | if err := (&ASTTool{}).applyRewrite([]astChange{{path: link, original: []byte("old"), content: []byte("new"), mode: 0o644}}); err != nil { | |
| 89 | + | t.Fatal(err) | |
| 90 | + | } | |
| 91 | + | if info, err := os.Lstat(link); err != nil || info.Mode()&os.ModeSymlink == 0 { | |
| 92 | + | t.Fatalf("link replaced: %v, %v", info, err) | |
| 93 | + | } | |
| 94 | + | if got, _ := os.ReadFile(target); string(got) != "new" { | |
| 95 | + | t.Fatalf("target = %q", got) | |
| 96 | + | } | |
| 97 | + | } | |
| 98 | + | ||
| 99 | + | func TestASTApplyRewriteRejectsSymlinkRetargetedOutsideJail(t *testing.T) { | |
| 100 | + | root := t.TempDir() | |
| 101 | + | outside := filepath.Join(t.TempDir(), "outside.go") | |
| 102 | + | link := filepath.Join(root, "link.go") | |
| 103 | + | if err := os.WriteFile(outside, []byte("old"), 0o644); err != nil { | |
| 104 | + | t.Fatal(err) | |
| 105 | + | } | |
| 106 | + | // The plan was made for an in-jail file, but the link now escapes. | |
| 107 | + | if err := os.Symlink(outside, link); err != nil { | |
| 108 | + | t.Skip("symlinks unavailable:", err) | |
| 109 | + | } | |
| 110 | + | sandbox := NewSandbox(root) | |
| 111 | + | sandbox.Lock() | |
| 112 | + | err := (&ASTTool{Sandbox: sandbox}).applyRewrite([]astChange{{path: link, original: []byte("old"), content: []byte("new"), mode: 0o644}}) | |
| 113 | + | if err == nil { | |
| 114 | + | t.Fatal("rewrite through escaping symlink succeeded") | |
| 115 | + | } | |
| 116 | + | if got, _ := os.ReadFile(outside); string(got) != "old" { | |
| 117 | + | t.Fatalf("outside file = %q", got) | |
| 118 | + | } | |
| 119 | + | } | |
| 120 | + | ||
| 121 | + | func TestASTGoCallRetryFailureIsReported(t *testing.T) { | |
| 122 | + | if runtime.GOOS == "windows" { | |
| 123 | + | t.Skip("test fixture uses a POSIX script") | |
| 124 | + | } | |
| 125 | + | root := t.TempDir() | |
| 126 | + | // `run` finds nothing (exit 1); the `scan` retry fails with an error. | |
| 127 | + | script := filepath.Join(root, "ast-grep") | |
| 128 | + | body := "#!/bin/sh\nif [ \"$1\" = scan ]; then echo 'Error: invalid rule' >&2; exit 8; fi\nexit 1\n" | |
| 129 | + | if err := os.WriteFile(script, []byte(body), 0o755); err != nil { | |
| 130 | + | t.Fatal(err) | |
| 131 | + | } | |
| 132 | + | tool := &ASTTool{CWD: root, LookPath: func(string) (string, error) { return script, nil }} | |
| 133 | + | _, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"fmt.Errorf($$$)","language":"go","path":"."}`), nil) | |
| 134 | + | if err == nil || !strings.Contains(err.Error(), "invalid rule") { | |
| 135 | + | t.Fatalf("Execute error = %v", err) | |
| 136 | + | } | |
| 137 | + | } | |
| 138 | + | ||
| 139 | + | func TestASTNoMatchIncludesPatternTip(t *testing.T) { | |
| 140 | + | root := t.TempDir() | |
| 141 | + | if err := os.WriteFile(filepath.Join(root, "main.go"), []byte("package main\n"), 0o644); err != nil { | |
| 142 | + | t.Fatal(err) | |
| 143 | + | } | |
| 144 | + | got := astText(t, realASTTool(t, root), `{"pattern":"missing($A)","language":"go","path":"."}`) | |
| 145 | + | if !strings.Contains(got, "No structural matches.") || !strings.Contains(got, "$$$A matches zero or more") { | |
| 146 | + | t.Fatalf("no-match text = %q", got) | |
| 147 | + | } | |
| 148 | + | } | |
| 149 | + | ||
| 150 | + | func TestASTGoCallPatternDetection(t *testing.T) { | |
| 151 | + | for _, tc := range []struct { | |
| 152 | + | args astArgs | |
| 153 | + | root string | |
| 154 | + | want bool | |
| 155 | + | }{ | |
| 156 | + | {astArgs{Pattern: "fmt.Errorf($$$)", Language: "go"}, "/src", true}, | |
| 157 | + | {astArgs{Pattern: "fmt.Errorf($$$)"}, "/src", true}, | |
| 158 | + | {astArgs{Pattern: "fmt.Errorf($$$)"}, "/src/main.go", true}, | |
| 159 | + | {astArgs{Pattern: "fmt.Errorf($$$)"}, "/src/app.ts", false}, | |
| 160 | + | {astArgs{Pattern: "console.log($$$)", Language: "typescript"}, "/src", false}, | |
| 161 | + | {astArgs{Pattern: "foo($A)", Language: "go"}, "/src", false}, | |
| 162 | + | {astArgs{Pattern: "if err != nil { $$$ }", Language: "go"}, "/src", false}, | |
| 163 | + | } { | |
| 164 | + | if got := astGoCallPattern(tc.args, tc.root); got != tc.want { | |
| 165 | + | t.Errorf("astGoCallPattern(%+v, %q) = %v, want %v", tc.args, tc.root, got, tc.want) | |
| 166 | + | } | |
| 167 | + | } | |
| 168 | + | } | |
| 169 | + | ||
| 170 | + | func TestASTSchemaAndDescriptionTeachUsage(t *testing.T) { | |
| 171 | + | tool := &ASTTool{} | |
| 172 | + | var schema struct { | |
| 173 | + | Required []string `json:"required"` | |
| 174 | + | } | |
| 175 | + | if err := json.Unmarshal(tool.Schema(), &schema); err != nil { | |
| 176 | + | t.Fatal(err) | |
| 177 | + | } | |
| 178 | + | if strings.Join(schema.Required, ",") != "pattern,path" { | |
| 179 | + | t.Fatalf("required = %v", schema.Required) | |
| 180 | + | } | |
| 181 | + | for _, want := range []string{"$$$", "rewrite", "grep"} { | |
| 182 | + | if !strings.Contains(tool.Description(), want) { | |
| 183 | + | t.Fatalf("description missing %q", want) | |
| 184 | + | } | |
| 185 | + | } | |
| 186 | + | } | |
| @@ -0,0 +1,186 @@ | |||||
| 1 | + | package tools | |||
| 2 | + | ||||
| 3 | + | import ( | |||
| 4 | + | "context" | |||
| 5 | + | "encoding/json" | |||
| 6 | + | "os" | |||
| 7 | + | "os/exec" | |||
| 8 | + | "path/filepath" | |||
| 9 | + | "runtime" | |||
| 10 | + | "strings" | |||
| 11 | + | "testing" | |||
| 12 | + | ||||
| 13 | + | "github.com/bnema/zut/packages/provider" | |||
| 14 | + | ) | |||
| 15 | + | ||||
| 16 | + | func realASTTool(t *testing.T, root string) *ASTTool { | |||
| 17 | + | t.Helper() | |||
| 18 | + | binary, err := exec.LookPath("ast-grep") | |||
| 19 | + | if err != nil { | |||
| 20 | + | t.Skip("ast-grep not installed") | |||
| 21 | + | } | |||
| 22 | + | return &ASTTool{CWD: root, LookPath: func(string) (string, error) { return binary, nil }} | |||
| 23 | + | } | |||
| 24 | + | ||||
| 25 | + | func astText(t *testing.T, tool *ASTTool, raw string) string { | |||
| 26 | + | t.Helper() | |||
| 27 | + | result, err := tool.Execute(context.Background(), json.RawMessage(raw), nil) | |||
| 28 | + | if err != nil { | |||
| 29 | + | t.Fatalf("Execute(%s): %v", raw, err) | |||
| 30 | + | } | |||
| 31 | + | return result.Content[0].(provider.TextBlock).Text | |||
| 32 | + | } | |||
| 33 | + | ||||
| 34 | + | func TestASTLanguageIsOptional(t *testing.T) { | |||
| 35 | + | root := t.TempDir() | |||
| 36 | + | if err := os.WriteFile(filepath.Join(root, "a.py"), []byte("print(1)\n"), 0o644); err != nil { | |||
| 37 | + | t.Fatal(err) | |||
| 38 | + | } | |||
| 39 | + | if err := os.WriteFile(filepath.Join(root, "b.ts"), []byte("console.log(2)\n"), 0o644); err != nil { | |||
| 40 | + | t.Fatal(err) | |||
| 41 | + | } | |||
| 42 | + | tool := realASTTool(t, root) | |||
| 43 | + | if got := astText(t, tool, `{"pattern":"print($A)","path":"."}`); !strings.Contains(got, "a.py:1") { | |||
| 44 | + | t.Fatalf("python search = %q", got) | |||
| 45 | + | } | |||
| 46 | + | if got := astText(t, tool, `{"pattern":"console.log($$$)","path":"."}`); !strings.Contains(got, "b.ts:1") { | |||
| 47 | + | t.Fatalf("typescript search = %q", got) | |||
| 48 | + | } | |||
| 49 | + | } | |||
| 50 | + | ||||
| 51 | + | // tree-sitter-go parses a bare `pkg.Func(args)` pattern as a type | |||
| 52 | + | // conversion, so without the retry these calls never match. | |||
| 53 | + | func TestASTGoQualifiedCallMatchesCalls(t *testing.T) { | |||
| 54 | + | root := t.TempDir() | |||
| 55 | + | path := filepath.Join(root, "main.go") | |||
| 56 | + | source := "package main\n\nimport \"fmt\"\n\nfunc main() {\n\t_ = fmt.Errorf(\"a\")\n\t_ = fmt.Errorf(\"b %d\",\n\t\t1)\n}\n" | |||
| 57 | + | if err := os.WriteFile(path, []byte(source), 0o644); err != nil { | |||
| 58 | + | t.Fatal(err) | |||
| 59 | + | } | |||
| 60 | + | tool := realASTTool(t, root) | |||
| 61 | + | for _, raw := range []string{ | |||
| 62 | + | `{"pattern":"fmt.Errorf($$$ARGS)","language":"go","path":"."}`, | |||
| 63 | + | `{"pattern":"fmt.Errorf($$$ARGS)","path":"."}`, | |||
| 64 | + | `{"pattern":"fmt.Errorf($$$ARGS)","path":"main.go"}`, | |||
| 65 | + | } { | |||
| 66 | + | got := astText(t, tool, raw) | |||
| 67 | + | if !strings.Contains(got, "main.go:6") || !strings.Contains(got, "main.go:7") { | |||
| 68 | + | t.Fatalf("search %s = %q", raw, got) | |||
| 69 | + | } | |||
| 70 | + | } | |||
| 71 | + | astText(t, tool, `{"pattern":"fmt.Errorf($$$ARGS)","path":".","rewrite":"errors.New($$$ARGS)"}`) | |||
| 72 | + | got, _ := os.ReadFile(path) | |||
| 73 | + | if strings.Contains(string(got), "fmt.Errorf") || strings.Count(string(got), "errors.New(") != 2 { | |||
| 74 | + | t.Fatalf("rewrite = %q", got) | |||
| 75 | + | } | |||
| 76 | + | } | |||
| 77 | + | ||||
| 78 | + | func TestASTRewriteWritesThroughSymlink(t *testing.T) { | |||
| 79 | + | root := t.TempDir() | |||
| 80 | + | target := filepath.Join(root, "real.go") | |||
| 81 | + | link := filepath.Join(root, "link.go") | |||
| 82 | + | if err := os.WriteFile(target, []byte("old"), 0o644); err != nil { | |||
| 83 | + | t.Fatal(err) | |||
| 84 | + | } | |||
| 85 | + | if err := os.Symlink(target, link); err != nil { | |||
| 86 | + | t.Skip("symlinks unavailable:", err) | |||
| 87 | + | } | |||
| 88 | + | if err := (&ASTTool{}).applyRewrite([]astChange{{path: link, original: []byte("old"), content: []byte("new"), mode: 0o644}}); err != nil { | |||
| 89 | + | t.Fatal(err) | |||
| 90 | + | } | |||
| 91 | + | if info, err := os.Lstat(link); err != nil || info.Mode()&os.ModeSymlink == 0 { | |||
| 92 | + | t.Fatalf("link replaced: %v, %v", info, err) | |||
| 93 | + | } | |||
| 94 | + | if got, _ := os.ReadFile(target); string(got) != "new" { | |||
| 95 | + | t.Fatalf("target = %q", got) | |||
| 96 | + | } | |||
| 97 | + | } | |||
| 98 | + | ||||
| 99 | + | func TestASTApplyRewriteRejectsSymlinkRetargetedOutsideJail(t *testing.T) { | |||
| 100 | + | root := t.TempDir() | |||
| 101 | + | outside := filepath.Join(t.TempDir(), "outside.go") | |||
| 102 | + | link := filepath.Join(root, "link.go") | |||
| 103 | + | if err := os.WriteFile(outside, []byte("old"), 0o644); err != nil { | |||
| 104 | + | t.Fatal(err) | |||
| 105 | + | } | |||
| 106 | + | // The plan was made for an in-jail file, but the link now escapes. | |||
| 107 | + | if err := os.Symlink(outside, link); err != nil { | |||
| 108 | + | t.Skip("symlinks unavailable:", err) | |||
| 109 | + | } | |||
| 110 | + | sandbox := NewSandbox(root) | |||
| 111 | + | sandbox.Lock() | |||
| 112 | + | err := (&ASTTool{Sandbox: sandbox}).applyRewrite([]astChange{{path: link, original: []byte("old"), content: []byte("new"), mode: 0o644}}) | |||
| 113 | + | if err == nil { | |||
| 114 | + | t.Fatal("rewrite through escaping symlink succeeded") | |||
| 115 | + | } | |||
| 116 | + | if got, _ := os.ReadFile(outside); string(got) != "old" { | |||
| 117 | + | t.Fatalf("outside file = %q", got) | |||
| 118 | + | } | |||
| 119 | + | } | |||
| 120 | + | ||||
| 121 | + | func TestASTGoCallRetryFailureIsReported(t *testing.T) { | |||
| 122 | + | if runtime.GOOS == "windows" { | |||
| 123 | + | t.Skip("test fixture uses a POSIX script") | |||
| 124 | + | } | |||
| 125 | + | root := t.TempDir() | |||
| 126 | + | // `run` finds nothing (exit 1); the `scan` retry fails with an error. | |||
| 127 | + | script := filepath.Join(root, "ast-grep") | |||
| 128 | + | body := "#!/bin/sh\nif [ \"$1\" = scan ]; then echo 'Error: invalid rule' >&2; exit 8; fi\nexit 1\n" | |||
| 129 | + | if err := os.WriteFile(script, []byte(body), 0o755); err != nil { | |||
| 130 | + | t.Fatal(err) | |||
| 131 | + | } | |||
| 132 | + | tool := &ASTTool{CWD: root, LookPath: func(string) (string, error) { return script, nil }} | |||
| 133 | + | _, err := tool.Execute(context.Background(), json.RawMessage(`{"pattern":"fmt.Errorf($$$)","language":"go","path":"."}`), nil) | |||
| 134 | + | if err == nil || !strings.Contains(err.Error(), "invalid rule") { | |||
| 135 | + | t.Fatalf("Execute error = %v", err) | |||
| 136 | + | } | |||
| 137 | + | } | |||
| 138 | + | ||||
| 139 | + | func TestASTNoMatchIncludesPatternTip(t *testing.T) { | |||
| 140 | + | root := t.TempDir() | |||
| 141 | + | if err := os.WriteFile(filepath.Join(root, "main.go"), []byte("package main\n"), 0o644); err != nil { | |||
| 142 | + | t.Fatal(err) | |||
| 143 | + | } | |||
| 144 | + | got := astText(t, realASTTool(t, root), `{"pattern":"missing($A)","language":"go","path":"."}`) | |||
| 145 | + | if !strings.Contains(got, "No structural matches.") || !strings.Contains(got, "$$$A matches zero or more") { | |||
| 146 | + | t.Fatalf("no-match text = %q", got) | |||
| 147 | + | } | |||
| 148 | + | } | |||
| 149 | + | ||||
| 150 | + | func TestASTGoCallPatternDetection(t *testing.T) { | |||
| 151 | + | for _, tc := range []struct { | |||
| 152 | + | args astArgs | |||
| 153 | + | root string | |||
| 154 | + | want bool | |||
| 155 | + | }{ | |||
| 156 | + | {astArgs{Pattern: "fmt.Errorf($$$)", Language: "go"}, "/src", true}, | |||
| 157 | + | {astArgs{Pattern: "fmt.Errorf($$$)"}, "/src", true}, | |||
| 158 | + | {astArgs{Pattern: "fmt.Errorf($$$)"}, "/src/main.go", true}, | |||
| 159 | + | {astArgs{Pattern: "fmt.Errorf($$$)"}, "/src/app.ts", false}, | |||
| 160 | + | {astArgs{Pattern: "console.log($$$)", Language: "typescript"}, "/src", false}, | |||
| 161 | + | {astArgs{Pattern: "foo($A)", Language: "go"}, "/src", false}, | |||
| 162 | + | {astArgs{Pattern: "if err != nil { $$$ }", Language: "go"}, "/src", false}, | |||
| 163 | + | } { | |||
| 164 | + | if got := astGoCallPattern(tc.args, tc.root); got != tc.want { | |||
| 165 | + | t.Errorf("astGoCallPattern(%+v, %q) = %v, want %v", tc.args, tc.root, got, tc.want) | |||
| 166 | + | } | |||
| 167 | + | } | |||
| 168 | + | } | |||
| 169 | + | ||||
| 170 | + | func TestASTSchemaAndDescriptionTeachUsage(t *testing.T) { | |||
| 171 | + | tool := &ASTTool{} | |||
| 172 | + | var schema struct { | |||
| 173 | + | Required []string `json:"required"` | |||
| 174 | + | } | |||
| 175 | + | if err := json.Unmarshal(tool.Schema(), &schema); err != nil { | |||
| 176 | + | t.Fatal(err) | |||
| 177 | + | } | |||
| 178 | + | if strings.Join(schema.Required, ",") != "pattern,path" { | |||
| 179 | + | t.Fatalf("required = %v", schema.Required) | |||
| 180 | + | } | |||
| 181 | + | for _, want := range []string{"$$$", "rewrite", "grep"} { | |||
| 182 | + | if !strings.Contains(tool.Description(), want) { | |||
| 183 | + | t.Fatalf("description missing %q", want) | |||
| 184 | + | } | |||
| 185 | + | } | |||
| 186 | + | } | |||
M packages/agent/tools/grep.go +41 -1
| @@ -9,6 +9,7 @@ import ( | |||
| 9 | 9 | "os" | |
| 10 | 10 | "os/exec" | |
| 11 | 11 | "path/filepath" | |
| 12 | + | "regexp" | |
| 12 | 13 | "strings" | |
| 13 | 14 | "sync" | |
| 14 | 15 | "time" | |
| @@ -33,6 +34,9 @@ type GrepTool struct { | |||
| 33 | 34 | CWD string | |
| 34 | 35 | Sandbox *Sandbox | |
| 35 | 36 | LookPath func(string) (string, error) | |
| 37 | + | // ASTAvailable reports whether the ast tool is registered and usable. | |
| 38 | + | // When nil or false, grep never suggests ast. | |
| 39 | + | ASTAvailable func() bool | |
| 36 | 40 | } | |
| 37 | 41 | ||
| 38 | 42 | var _ core.Tool = (*GrepTool)(nil) | |
| @@ -176,10 +180,21 @@ func (t *GrepTool) executeWithTimeout(ctx context.Context, raw json.RawMessage, | |||
| 176 | 180 | // diagnostics while reporting the normalized tool exit code above. | |
| 177 | 181 | isError = false | |
| 178 | 182 | } | |
| 183 | + | ||
| 184 | + | discovers := grepResultPaths(text, root, info.IsDir()) | |
| 185 | + | // The hint is appended after path discovery so it is never read as a | |
| 186 | + | // "file:line" result. | |
| 187 | + | hint := "" | |
| 188 | + | if !isError && t.ASTAvailable != nil { | |
| 189 | + | if hint = grepASTHint(args.Pattern); hint != "" && !t.ASTAvailable() { | |
| 190 | + | hint = "" | |
| 191 | + | } | |
| 192 | + | } | |
| 193 | + | text += hint | |
| 179 | 194 | return core.ToolResult{ | |
| 180 | 195 | Content: []provider.Content{provider.TextBlock{Text: text}}, | |
| 181 | 196 | IsError: isError, | |
| 182 | - | Context: provider.ToolContext{Discovers: grepResultPaths(text, root, info.IsDir())}, | |
| 197 | + | Context: provider.ToolContext{Discovers: discovers}, | |
| 183 | 198 | Details: map[string]any{ | |
| 184 | 199 | "engine": engine, | |
| 185 | 200 | "path": t.Sandbox.DisplayPath(root, args.Path), | |
| @@ -189,10 +204,35 @@ func (t *GrepTool) executeWithTimeout(ctx context.Context, raw json.RawMessage, | |||
| 189 | 204 | "bytes_truncated": output.bytesTruncated(), | |
| 190 | 205 | "lines_truncated": linesTruncated, | |
| 191 | 206 | "output_cancelled": outputCancelled, | |
| 207 | + | "ast_hint": hint != "", | |
| 192 | 208 | }, | |
| 193 | 209 | }, nil | |
| 194 | 210 | } | |
| 195 | 211 | ||
| 212 | + | var ( | |
| 213 | + | // grepCallPattern matches regexes that search for a call, such as | |
| 214 | + | // `fmt\.Errorf\(` or `\bNewClient\(`. The capture is the callee. | |
| 215 | + | // Qualified names need an escaped dot: an unescaped `.` matches any | |
| 216 | + | // character, so a literal structural pattern would be narrower. | |
| 217 | + | grepCallPattern = regexp.MustCompile(`^(?:\\b)?([A-Za-z_]\w*(?:\\\.[A-Za-z_]\w*)*)\\\($`) | |
| 218 | + | // grepDefinitionPattern matches regexes that search for a definition. | |
| 219 | + | grepDefinitionPattern = regexp.MustCompile(`^\^?(?:\\s\*|\s)*(?:func|def|class|fn|function|interface|struct|impl)(?:\s|\\s|\\b|\\\()`) | |
| 220 | + | ) | |
| 221 | + | ||
| 222 | + | // grepASTHint suggests the structural ast tool when a regex is really a | |
| 223 | + | // search for code shape. Regexes miss calls split across lines and cannot | |
| 224 | + | // tell code from comments or strings; ast-grep handles both. | |
| 225 | + | func grepASTHint(pattern string) string { | |
| 226 | + | if match := grepCallPattern.FindStringSubmatch(pattern); match != nil { | |
| 227 | + | callee := strings.ReplaceAll(match[1], `\.`, ".") | |
| 228 | + | return fmt.Sprintf("\nTip: for code structure, the ast tool matches calls across lines and skips comments and strings: ast {\"pattern\":%q,\"path\":\"...\"}. Add \"rewrite\" to change every call at once.\n", callee+"($$$ARGS)") | |
| 229 | + | } | |
| 230 | + | if grepDefinitionPattern.MatchString(pattern) { | |
| 231 | + | return "\nTip: for definitions, the ast tool matches code shape, e.g. ast {\"pattern\":\"func $NAME($$$) $$$ { $$$ }\",\"path\":\"...\"}.\n" | |
| 232 | + | } | |
| 233 | + | return "" | |
| 234 | + | } | |
| 235 | + | ||
| 196 | 236 | func grepResultPaths(output, root string, rootIsDir bool) []string { | |
| 197 | 237 | seen := make(map[string]bool) | |
| 198 | 238 | var paths []string | |
| @@ -9,6 +9,7 @@ import ( | |||||
| 9 | "os" | 9 | "os" | ||
| 10 | "os/exec" | 10 | "os/exec" | ||
| 11 | "path/filepath" | 11 | "path/filepath" | ||
| 12 | + | "regexp" | |||
| 12 | "strings" | 13 | "strings" | ||
| 13 | "sync" | 14 | "sync" | ||
| 14 | "time" | 15 | "time" | ||
| @@ -33,6 +34,9 @@ type GrepTool struct { | |||||
| 33 | CWD string | 34 | CWD string | ||
| 34 | Sandbox *Sandbox | 35 | Sandbox *Sandbox | ||
| 35 | LookPath func(string) (string, error) | 36 | LookPath func(string) (string, error) | ||
| 37 | + | // ASTAvailable reports whether the ast tool is registered and usable. | |||
| 38 | + | // When nil or false, grep never suggests ast. | |||
| 39 | + | ASTAvailable func() bool | |||
| 36 | } | 40 | } | ||
| 37 | 41 | ||||
| 38 | var _ core.Tool = (*GrepTool)(nil) | 42 | var _ core.Tool = (*GrepTool)(nil) | ||
| @@ -176,10 +180,21 @@ func (t *GrepTool) executeWithTimeout(ctx context.Context, raw json.RawMessage, | |||||
| 176 | // diagnostics while reporting the normalized tool exit code above. | 180 | // diagnostics while reporting the normalized tool exit code above. | ||
| 177 | isError = false | 181 | isError = false | ||
| 178 | } | 182 | } | ||
| 183 | + | ||||
| 184 | + | discovers := grepResultPaths(text, root, info.IsDir()) | |||
| 185 | + | // The hint is appended after path discovery so it is never read as a | |||
| 186 | + | // "file:line" result. | |||
| 187 | + | hint := "" | |||
| 188 | + | if !isError && t.ASTAvailable != nil { | |||
| 189 | + | if hint = grepASTHint(args.Pattern); hint != "" && !t.ASTAvailable() { | |||
| 190 | + | hint = "" | |||
| 191 | + | } | |||
| 192 | + | } | |||
| 193 | + | text += hint | |||
| 179 | return core.ToolResult{ | 194 | return core.ToolResult{ | ||
| 180 | Content: []provider.Content{provider.TextBlock{Text: text}}, | 195 | Content: []provider.Content{provider.TextBlock{Text: text}}, | ||
| 181 | IsError: isError, | 196 | IsError: isError, | ||
| 182 | - | Context: provider.ToolContext{Discovers: grepResultPaths(text, root, info.IsDir())}, | 197 | + | Context: provider.ToolContext{Discovers: discovers}, |
| 183 | Details: map[string]any{ | 198 | Details: map[string]any{ | ||
| 184 | "engine": engine, | 199 | "engine": engine, | ||
| 185 | "path": t.Sandbox.DisplayPath(root, args.Path), | 200 | "path": t.Sandbox.DisplayPath(root, args.Path), | ||
| @@ -189,10 +204,35 @@ func (t *GrepTool) executeWithTimeout(ctx context.Context, raw json.RawMessage, | |||||
| 189 | "bytes_truncated": output.bytesTruncated(), | 204 | "bytes_truncated": output.bytesTruncated(), | ||
| 190 | "lines_truncated": linesTruncated, | 205 | "lines_truncated": linesTruncated, | ||
| 191 | "output_cancelled": outputCancelled, | 206 | "output_cancelled": outputCancelled, | ||
| 207 | + | "ast_hint": hint != "", | |||
| 192 | }, | 208 | }, | ||
| 193 | }, nil | 209 | }, nil | ||
| 194 | } | 210 | } | ||
| 195 | 211 | ||||
| 212 | + | var ( | |||
| 213 | + | // grepCallPattern matches regexes that search for a call, such as | |||
| 214 | + | // `fmt\.Errorf\(` or `\bNewClient\(`. The capture is the callee. | |||
| 215 | + | // Qualified names need an escaped dot: an unescaped `.` matches any | |||
| 216 | + | // character, so a literal structural pattern would be narrower. | |||
| 217 | + | grepCallPattern = regexp.MustCompile(`^(?:\\b)?([A-Za-z_]\w*(?:\\\.[A-Za-z_]\w*)*)\\\($`) | |||
| 218 | + | // grepDefinitionPattern matches regexes that search for a definition. | |||
| 219 | + | grepDefinitionPattern = regexp.MustCompile(`^\^?(?:\\s\*|\s)*(?:func|def|class|fn|function|interface|struct|impl)(?:\s|\\s|\\b|\\\()`) | |||
| 220 | + | ) | |||
| 221 | + | ||||
| 222 | + | // grepASTHint suggests the structural ast tool when a regex is really a | |||
| 223 | + | // search for code shape. Regexes miss calls split across lines and cannot | |||
| 224 | + | // tell code from comments or strings; ast-grep handles both. | |||
| 225 | + | func grepASTHint(pattern string) string { | |||
| 226 | + | if match := grepCallPattern.FindStringSubmatch(pattern); match != nil { | |||
| 227 | + | callee := strings.ReplaceAll(match[1], `\.`, ".") | |||
| 228 | + | return fmt.Sprintf("\nTip: for code structure, the ast tool matches calls across lines and skips comments and strings: ast {\"pattern\":%q,\"path\":\"...\"}. Add \"rewrite\" to change every call at once.\n", callee+"($$$ARGS)") | |||
| 229 | + | } | |||
| 230 | + | if grepDefinitionPattern.MatchString(pattern) { | |||
| 231 | + | return "\nTip: for definitions, the ast tool matches code shape, e.g. ast {\"pattern\":\"func $NAME($$$) $$$ { $$$ }\",\"path\":\"...\"}.\n" | |||
| 232 | + | } | |||
| 233 | + | return "" | |||
| 234 | + | } | |||
| 235 | + | ||||
| 196 | func grepResultPaths(output, root string, rootIsDir bool) []string { | 236 | func grepResultPaths(output, root string, rootIsDir bool) []string { | ||
| 197 | seen := make(map[string]bool) | 237 | seen := make(map[string]bool) | ||
| 198 | var paths []string | 238 | var paths []string | ||
A packages/agent/tools/grep_ast_hint_test.go +73 -0
| @@ -0,0 +1,73 @@ | |||
| 1 | + | package tools | |
| 2 | + | ||
| 3 | + | import ( | |
| 4 | + | "context" | |
| 5 | + | "encoding/json" | |
| 6 | + | "path/filepath" | |
| 7 | + | "strings" | |
| 8 | + | "testing" | |
| 9 | + | ) | |
| 10 | + | ||
| 11 | + | func TestGrepASTHintClassifiesPatterns(t *testing.T) { | |
| 12 | + | for _, tc := range []struct { | |
| 13 | + | pattern string | |
| 14 | + | want string // substring of the hint; empty means no hint | |
| 15 | + | }{ | |
| 16 | + | {`fmt\.Errorf\(`, `"fmt.Errorf($$$ARGS)"`}, | |
| 17 | + | {`\bNewClient\(`, `"NewClient($$$ARGS)"`}, | |
| 18 | + | {`os\.ReadFile\(`, `"os.ReadFile($$$ARGS)"`}, | |
| 19 | + | {`os.ReadFile\(`, ""}, | |
| 20 | + | {`^func `, "for definitions"}, | |
| 21 | + | {`^\s*def\s`, "for definitions"}, | |
| 22 | + | {`class\b`, "for definitions"}, | |
| 23 | + | {`TODO`, ""}, | |
| 24 | + | {`fmt\.Errorf`, ""}, | |
| 25 | + | {`error: .*\(`, ""}, | |
| 26 | + | {`foo\(bar\)`, ""}, | |
| 27 | + | {`functional`, ""}, | |
| 28 | + | } { | |
| 29 | + | got := grepASTHint(tc.pattern) | |
| 30 | + | if tc.want == "" { | |
| 31 | + | if got != "" { | |
| 32 | + | t.Errorf("grepASTHint(%q) = %q, want none", tc.pattern, got) | |
| 33 | + | } | |
| 34 | + | continue | |
| 35 | + | } | |
| 36 | + | if !strings.Contains(got, tc.want) { | |
| 37 | + | t.Errorf("grepASTHint(%q) = %q, want substring %q", tc.pattern, got, tc.want) | |
| 38 | + | } | |
| 39 | + | } | |
| 40 | + | } | |
| 41 | + | ||
| 42 | + | func TestGrepAppendsASTHintOnlyWhenAvailable(t *testing.T) { | |
| 43 | + | skipGrepScriptTestsOnWindows(t) | |
| 44 | + | root := t.TempDir() | |
| 45 | + | rg := writeGrepScript(t, "#!/bin/sh\nprintf 'main.go:3:\\treturn fmt.Errorf(\"x\")\\n'\n") | |
| 46 | + | args := json.RawMessage(`{"pattern":"fmt\\.Errorf\\(","path":"."}`) | |
| 47 | + | for _, tc := range []struct { | |
| 48 | + | name string | |
| 49 | + | available func() bool | |
| 50 | + | wantHint bool | |
| 51 | + | }{ | |
| 52 | + | {"not wired", nil, false}, | |
| 53 | + | {"missing backend", func() bool { return false }, false}, | |
| 54 | + | {"available", func() bool { return true }, true}, | |
| 55 | + | } { | |
| 56 | + | t.Run(tc.name, func(t *testing.T) { | |
| 57 | + | tool := &GrepTool{CWD: root, ASTAvailable: tc.available, LookPath: lookupScripts(map[string]string{"rg": rg})} | |
| 58 | + | result, err := tool.Execute(context.Background(), args, nil) | |
| 59 | + | if err != nil { | |
| 60 | + | t.Fatal(err) | |
| 61 | + | } | |
| 62 | + | text := grepText(result) | |
| 63 | + | if got := strings.Contains(text, "Tip:"); got != tc.wantHint { | |
| 64 | + | t.Fatalf("hint present = %v, want %v; output %q", got, tc.wantHint, text) | |
| 65 | + | } | |
| 66 | + | // The hint must never be read back as a discovered file. | |
| 67 | + | want := filepath.Join(root, "main.go") | |
| 68 | + | if len(result.Context.Discovers) != 1 || result.Context.Discovers[0] != want { | |
| 69 | + | t.Fatalf("discovers = %v, want [%s]", result.Context.Discovers, want) | |
| 70 | + | } | |
| 71 | + | }) | |
| 72 | + | } | |
| 73 | + | } | |
| @@ -0,0 +1,73 @@ | |||||
| 1 | + | package tools | |||
| 2 | + | ||||
| 3 | + | import ( | |||
| 4 | + | "context" | |||
| 5 | + | "encoding/json" | |||
| 6 | + | "path/filepath" | |||
| 7 | + | "strings" | |||
| 8 | + | "testing" | |||
| 9 | + | ) | |||
| 10 | + | ||||
| 11 | + | func TestGrepASTHintClassifiesPatterns(t *testing.T) { | |||
| 12 | + | for _, tc := range []struct { | |||
| 13 | + | pattern string | |||
| 14 | + | want string // substring of the hint; empty means no hint | |||
| 15 | + | }{ | |||
| 16 | + | {`fmt\.Errorf\(`, `"fmt.Errorf($$$ARGS)"`}, | |||
| 17 | + | {`\bNewClient\(`, `"NewClient($$$ARGS)"`}, | |||
| 18 | + | {`os\.ReadFile\(`, `"os.ReadFile($$$ARGS)"`}, | |||
| 19 | + | {`os.ReadFile\(`, ""}, | |||
| 20 | + | {`^func `, "for definitions"}, | |||
| 21 | + | {`^\s*def\s`, "for definitions"}, | |||
| 22 | + | {`class\b`, "for definitions"}, | |||
| 23 | + | {`TODO`, ""}, | |||
| 24 | + | {`fmt\.Errorf`, ""}, | |||
| 25 | + | {`error: .*\(`, ""}, | |||
| 26 | + | {`foo\(bar\)`, ""}, | |||
| 27 | + | {`functional`, ""}, | |||
| 28 | + | } { | |||
| 29 | + | got := grepASTHint(tc.pattern) | |||
| 30 | + | if tc.want == "" { | |||
| 31 | + | if got != "" { | |||
| 32 | + | t.Errorf("grepASTHint(%q) = %q, want none", tc.pattern, got) | |||
| 33 | + | } | |||
| 34 | + | continue | |||
| 35 | + | } | |||
| 36 | + | if !strings.Contains(got, tc.want) { | |||
| 37 | + | t.Errorf("grepASTHint(%q) = %q, want substring %q", tc.pattern, got, tc.want) | |||
| 38 | + | } | |||
| 39 | + | } | |||
| 40 | + | } | |||
| 41 | + | ||||
| 42 | + | func TestGrepAppendsASTHintOnlyWhenAvailable(t *testing.T) { | |||
| 43 | + | skipGrepScriptTestsOnWindows(t) | |||
| 44 | + | root := t.TempDir() | |||
| 45 | + | rg := writeGrepScript(t, "#!/bin/sh\nprintf 'main.go:3:\\treturn fmt.Errorf(\"x\")\\n'\n") | |||
| 46 | + | args := json.RawMessage(`{"pattern":"fmt\\.Errorf\\(","path":"."}`) | |||
| 47 | + | for _, tc := range []struct { | |||
| 48 | + | name string | |||
| 49 | + | available func() bool | |||
| 50 | + | wantHint bool | |||
| 51 | + | }{ | |||
| 52 | + | {"not wired", nil, false}, | |||
| 53 | + | {"missing backend", func() bool { return false }, false}, | |||
| 54 | + | {"available", func() bool { return true }, true}, | |||
| 55 | + | } { | |||
| 56 | + | t.Run(tc.name, func(t *testing.T) { | |||
| 57 | + | tool := &GrepTool{CWD: root, ASTAvailable: tc.available, LookPath: lookupScripts(map[string]string{"rg": rg})} | |||
| 58 | + | result, err := tool.Execute(context.Background(), args, nil) | |||
| 59 | + | if err != nil { | |||
| 60 | + | t.Fatal(err) | |||
| 61 | + | } | |||
| 62 | + | text := grepText(result) | |||
| 63 | + | if got := strings.Contains(text, "Tip:"); got != tc.wantHint { | |||
| 64 | + | t.Fatalf("hint present = %v, want %v; output %q", got, tc.wantHint, text) | |||
| 65 | + | } | |||
| 66 | + | // The hint must never be read back as a discovered file. | |||
| 67 | + | want := filepath.Join(root, "main.go") | |||
| 68 | + | if len(result.Context.Discovers) != 1 || result.Context.Discovers[0] != want { | |||
| 69 | + | t.Fatalf("discovers = %v, want [%s]", result.Context.Discovers, want) | |||
| 70 | + | } | |||
| 71 | + | }) | |||
| 72 | + | } | |||
| 73 | + | } | |||