|
5 | 5 | import sys
|
6 | 6 | from typing import Any, List, Literal, Optional
|
7 | 7 |
|
| 8 | +import yaml |
8 | 9 | from mcp.server.fastmcp import FastMCP
|
9 | 10 | from pydantic import Field
|
10 | 11 |
|
@@ -59,168 +60,171 @@ def parse_args_and_get_config():
|
59 | 60 |
|
60 | 61 | DumpFormat = Literal["pattern", "cst", "ast"]
|
61 | 62 |
|
62 |
| -@mcp.tool() |
63 |
| -def dump_syntax_tree( |
64 |
| - code: str = Field(description = "The code you need"), |
65 |
| - language: str = Field(description = "The language of the code"), |
66 |
| - format: DumpFormat = Field(description = "Code dump format. Available values: pattern, ast, cst", default = "cst"), |
67 |
| -) -> str: |
68 |
| - """ |
69 |
| - Dump code's syntax structure or dump a query's pattern structure. |
70 |
| - This is useful to discover correct syntax kind and syntax tree structure. Call it when debugging a rule. |
71 |
| - The tool requires three arguments: code, language and format. The first two are self-explanatory. |
72 |
| - `format` is the output format of the syntax tree. |
73 |
| - use `format=cst` to inspect the code's concrete syntax tree structure, useful to debug target code. |
74 |
| - use `format=pattern` to inspect how ast-grep interprets a pattern, useful to debug pattern rule. |
75 |
| -
|
76 |
| - Internally calls: ast-grep run --pattern <code> --lang <language> --debug-query=<format> |
77 |
| - """ |
78 |
| - result = run_ast_grep("run", ["--pattern", code, "--lang", language, f"--debug-query={format}"]) |
79 |
| - return result.stderr.strip() # type: ignore[no-any-return] |
80 |
| - |
81 |
| -@mcp.tool() |
82 |
| -def test_match_code_rule( |
83 |
| - code: str = Field(description="The code to test against the rule"), |
84 |
| - yaml: str = Field(description="The ast-grep YAML rule to search. It must have id, language, rule fields."), |
85 |
| -) -> List[dict[str, Any]]: |
86 |
| - """ |
87 |
| - Test a code against an ast-grep YAML rule. |
88 |
| - This is useful to test a rule before using it in a project. |
89 |
| -
|
90 |
| - Internally calls: ast-grep scan --inline-rules <yaml> --json --stdin |
91 |
| - """ |
92 |
| - result = run_ast_grep("scan", ["--inline-rules", yaml, "--json", "--stdin"], input_text = code) |
93 |
| - matches = json.loads(result.stdout.strip()) |
94 |
| - if not matches: |
95 |
| - raise ValueError("No matches found for the given code and rule. Try adding `stopBy: end` to your inside/has rule.") |
96 |
| - return matches # type: ignore[no-any-return] |
97 |
| - |
98 |
| -@mcp.tool() |
99 |
| -def find_code( |
100 |
| - project_folder: str = Field(description="The absolute path to the project folder. It must be absolute path."), |
101 |
| - pattern: str = Field(description="The ast-grep pattern to search for. Note, the pattern must have valid AST structure."), |
102 |
| - language: str = Field(description="The language of the query", default=""), |
103 |
| - max_results: Optional[int] = Field(default=None, description="Maximum results to return"), |
104 |
| - output_format: str = Field(default="text", description="'text' or 'json'"), |
105 |
| -) -> str | List[dict[str, Any]]: |
106 |
| - """ |
107 |
| - Find code in a project folder that matches the given ast-grep pattern. |
108 |
| - Pattern is good for simple and single-AST node result. |
109 |
| - For more complex usage, please use YAML by `find_code_by_rule`. |
110 |
| -
|
111 |
| - Internally calls: ast-grep run --pattern <pattern> [--json] <project_folder> |
112 |
| -
|
113 |
| - Output formats: |
114 |
| - - text (default): Compact text format with file:line-range headers and complete match text |
115 |
| - Example: |
116 |
| - Found 2 matches: |
117 |
| -
|
118 |
| - path/to/file.py:10-15 |
119 |
| - def example_function(): |
120 |
| - # function body |
121 |
| - return result |
122 |
| -
|
123 |
| - path/to/file.py:20-22 |
124 |
| - def another_function(): |
125 |
| - pass |
126 |
| -
|
127 |
| - - json: Full match objects with metadata including ranges, meta-variables, etc. |
128 |
| -
|
129 |
| - The max_results parameter limits the number of complete matches returned (not individual lines). |
130 |
| - When limited, the header shows "Found X matches (showing first Y of Z)". |
131 |
| -
|
132 |
| - Example usage: |
133 |
| - find_code(pattern="class $NAME", max_results=20) # Returns text format |
134 |
| - find_code(pattern="class $NAME", output_format="json") # Returns JSON with metadata |
135 |
| - """ |
136 |
| - if output_format not in ["text", "json"]: |
137 |
| - raise ValueError(f"Invalid output_format: {output_format}. Must be 'text' or 'json'.") |
138 |
| - |
139 |
| - args = ["--pattern", pattern] |
140 |
| - if language: |
141 |
| - args.extend(["--lang", language]) |
142 |
| - |
143 |
| - # Always get JSON internally for accurate match limiting |
144 |
| - result = run_ast_grep("run", args + ["--json", project_folder]) |
145 |
| - matches = json.loads(result.stdout.strip() or "[]") |
146 |
| - |
147 |
| - # Apply max_results limit to complete matches |
148 |
| - total_matches = len(matches) |
149 |
| - if max_results is not None and total_matches > max_results: |
150 |
| - matches = matches[:max_results] |
151 |
| - |
152 |
| - if output_format == "text": |
| 63 | +def register_mcp_tools() -> None: |
| 64 | + @mcp.tool() |
| 65 | + def dump_syntax_tree( |
| 66 | + code: str = Field(description = "The code you need"), |
| 67 | + language: str = Field(description = f"The language of the code. Supported: {', '.join(get_supported_languages())}"), |
| 68 | + format: DumpFormat = Field(description = "Code dump format. Available values: pattern, ast, cst", default = "cst"), |
| 69 | + ) -> str: |
| 70 | + """ |
| 71 | + Dump code's syntax structure or dump a query's pattern structure. |
| 72 | + This is useful to discover correct syntax kind and syntax tree structure. Call it when debugging a rule. |
| 73 | + The tool requires three arguments: code, language and format. The first two are self-explanatory. |
| 74 | + `format` is the output format of the syntax tree. |
| 75 | + use `format=cst` to inspect the code's concrete syntax tree structure, useful to debug target code. |
| 76 | + use `format=pattern` to inspect how ast-grep interprets a pattern, useful to debug pattern rule. |
| 77 | +
|
| 78 | + Internally calls: ast-grep run --pattern <code> --lang <language> --debug-query=<format> |
| 79 | + """ |
| 80 | + result = run_ast_grep("run", ["--pattern", code, "--lang", language, f"--debug-query={format}"]) |
| 81 | + return result.stderr.strip() # type: ignore[no-any-return] |
| 82 | + |
| 83 | + @mcp.tool() |
| 84 | + def test_match_code_rule( |
| 85 | + code: str = Field(description = "The code to test against the rule"), |
| 86 | + yaml: str = Field(description = "The ast-grep YAML rule to search. It must have id, language, rule fields."), |
| 87 | + ) -> List[dict[str, Any]]: |
| 88 | + """ |
| 89 | + Test a code against an ast-grep YAML rule. |
| 90 | + This is useful to test a rule before using it in a project. |
| 91 | +
|
| 92 | + Internally calls: ast-grep scan --inline-rules <yaml> --json --stdin |
| 93 | + """ |
| 94 | + result = run_ast_grep("scan", ["--inline-rules", yaml, "--json", "--stdin"], input_text = code) |
| 95 | + matches = json.loads(result.stdout.strip()) |
153 | 96 | if not matches:
|
154 |
| - return "No matches found" |
155 |
| - text_output = format_matches_as_text(matches) |
156 |
| - header = f"Found {len(matches)} matches" |
157 |
| - if max_results is not None and total_matches > max_results: |
158 |
| - header += f" (showing first {max_results} of {total_matches})" |
159 |
| - return header + ":\n\n" + text_output |
160 |
| - return matches # type: ignore[no-any-return] |
161 |
| - |
162 |
| -@mcp.tool() |
163 |
| -def find_code_by_rule( |
164 |
| - project_folder: str = Field(description="The absolute path to the project folder. It must be absolute path."), |
165 |
| - yaml: str = Field(description="The ast-grep YAML rule to search. It must have id, language, rule fields."), |
166 |
| - max_results: Optional[int] = Field(default=None, description="Maximum results to return"), |
167 |
| - output_format: str = Field(default="text", description="'text' or 'json'"), |
| 97 | + raise ValueError("No matches found for the given code and rule. Try adding `stopBy: end` to your inside/has rule.") |
| 98 | + return matches # type: ignore[no-any-return] |
| 99 | + |
| 100 | + @mcp.tool() |
| 101 | + def find_code( |
| 102 | + project_folder: str = Field(description = "The absolute path to the project folder. It must be absolute path."), |
| 103 | + pattern: str = Field(description = "The ast-grep pattern to search for. Note, the pattern must have valid AST structure."), |
| 104 | + language: str = Field(description = f"The language of the code. Supported: {', '.join(get_supported_languages())}. " |
| 105 | + "If not specified, will be auto-detected based on file extensions.", default = ""), |
| 106 | + max_results: Optional[int] = Field(default = None, description = "Maximum results to return"), |
| 107 | + output_format: str = Field(default = "text", description = "'text' or 'json'"), |
168 | 108 | ) -> str | List[dict[str, Any]]:
|
169 |
| - """ |
170 |
| - Find code using ast-grep's YAML rule in a project folder. |
171 |
| - YAML rule is more powerful than simple pattern and can perform complex search like find AST inside/having another AST. |
172 |
| - It is a more advanced search tool than the simple `find_code`. |
| 109 | + """ |
| 110 | + Find code in a project folder that matches the given ast-grep pattern. |
| 111 | + Pattern is good for simple and single-AST node result. |
| 112 | + For more complex usage, please use YAML by `find_code_by_rule`. |
173 | 113 |
|
174 |
| - Tip: When using relational rules (inside/has), add `stopBy: end` to ensure complete traversal. |
| 114 | + Internally calls: ast-grep run --pattern <pattern> [--json] <project_folder> |
175 | 115 |
|
176 |
| - Internally calls: ast-grep scan --inline-rules <yaml> [--json] <project_folder> |
| 116 | + Output formats: |
| 117 | + - text (default): Compact text format with file:line-range headers and complete match text |
| 118 | + Example: |
| 119 | + Found 2 matches: |
177 | 120 |
|
178 |
| - Output formats: |
179 |
| - - text (default): Compact text format with file:line-range headers and complete match text |
180 |
| - Example: |
181 |
| - Found 2 matches: |
| 121 | + path/to/file.py:10-15 |
| 122 | + def example_function(): |
| 123 | + # function body |
| 124 | + return result |
182 | 125 |
|
183 |
| - src/models.py:45-52 |
184 |
| - class UserModel: |
185 |
| - def __init__(self): |
186 |
| - self.id = None |
187 |
| - self.name = None |
| 126 | + path/to/file.py:20-22 |
| 127 | + def another_function(): |
| 128 | + pass |
188 | 129 |
|
189 |
| - src/views.py:12 |
190 |
| - class SimpleView: pass |
| 130 | + - json: Full match objects with metadata including ranges, meta-variables, etc. |
191 | 131 |
|
192 |
| - - json: Full match objects with metadata including ranges, meta-variables, etc. |
| 132 | + The max_results parameter limits the number of complete matches returned (not individual lines). |
| 133 | + When limited, the header shows "Found X matches (showing first Y of Z)". |
193 | 134 |
|
194 |
| - The max_results parameter limits the number of complete matches returned (not individual lines). |
195 |
| - When limited, the header shows "Found X matches (showing first Y of Z)". |
| 135 | + Example usage: |
| 136 | + find_code(pattern="class $NAME", max_results=20) # Returns text format |
| 137 | + find_code(pattern="class $NAME", output_format="json") # Returns JSON with metadata |
| 138 | + """ |
| 139 | + if output_format not in ["text", "json"]: |
| 140 | + raise ValueError(f"Invalid output_format: {output_format}. Must be 'text' or 'json'.") |
196 | 141 |
|
197 |
| - Example usage: |
198 |
| - find_code_by_rule(yaml="id: x\\nlanguage: python\\nrule: {pattern: 'class $NAME'}", max_results=20) |
199 |
| - find_code_by_rule(yaml="...", output_format="json") # For full metadata |
200 |
| - """ |
201 |
| - if output_format not in ["text", "json"]: |
202 |
| - raise ValueError(f"Invalid output_format: {output_format}. Must be 'text' or 'json'.") |
| 142 | + args = ["--pattern", pattern] |
| 143 | + if language: |
| 144 | + args.extend(["--lang", language]) |
203 | 145 |
|
204 |
| - args = ["--inline-rules", yaml] |
| 146 | + # Always get JSON internally for accurate match limiting |
| 147 | + result = run_ast_grep("run", args + ["--json", project_folder]) |
| 148 | + matches = json.loads(result.stdout.strip() or "[]") |
205 | 149 |
|
206 |
| - # Always get JSON internally for accurate match limiting |
207 |
| - result = run_ast_grep("scan", args + ["--json", project_folder]) |
208 |
| - matches = json.loads(result.stdout.strip() or "[]") |
| 150 | + # Apply max_results limit to complete matches |
| 151 | + total_matches = len(matches) |
| 152 | + if max_results is not None and total_matches > max_results: |
| 153 | + matches = matches[:max_results] |
| 154 | + |
| 155 | + if output_format == "text": |
| 156 | + if not matches: |
| 157 | + return "No matches found" |
| 158 | + text_output = format_matches_as_text(matches) |
| 159 | + header = f"Found {len(matches)} matches" |
| 160 | + if max_results is not None and total_matches > max_results: |
| 161 | + header += f" (showing first {max_results} of {total_matches})" |
| 162 | + return header + ":\n\n" + text_output |
| 163 | + return matches # type: ignore[no-any-return] |
| 164 | + |
| 165 | + @mcp.tool() |
| 166 | + def find_code_by_rule( |
| 167 | + project_folder: str = Field(description = "The absolute path to the project folder. It must be absolute path."), |
| 168 | + yaml: str = Field(description = "The ast-grep YAML rule to search. It must have id, language, rule fields."), |
| 169 | + max_results: Optional[int] = Field(default = None, description = "Maximum results to return"), |
| 170 | + output_format: str = Field(default = "text", description = "'text' or 'json'"), |
| 171 | + ) -> str | List[dict[str, Any]]: |
| 172 | + """ |
| 173 | + Find code using ast-grep's YAML rule in a project folder. |
| 174 | + YAML rule is more powerful than simple pattern and can perform complex search like find AST inside/having another AST. |
| 175 | + It is a more advanced search tool than the simple `find_code`. |
| 176 | +
|
| 177 | + Tip: When using relational rules (inside/has), add `stopBy: end` to ensure complete traversal. |
| 178 | +
|
| 179 | + Internally calls: ast-grep scan --inline-rules <yaml> [--json] <project_folder> |
| 180 | +
|
| 181 | + Output formats: |
| 182 | + - text (default): Compact text format with file:line-range headers and complete match text |
| 183 | + Example: |
| 184 | + Found 2 matches: |
| 185 | +
|
| 186 | + src/models.py:45-52 |
| 187 | + class UserModel: |
| 188 | + def __init__(self): |
| 189 | + self.id = None |
| 190 | + self.name = None |
| 191 | +
|
| 192 | + src/views.py:12 |
| 193 | + class SimpleView: pass |
| 194 | +
|
| 195 | + - json: Full match objects with metadata including ranges, meta-variables, etc. |
| 196 | +
|
| 197 | + The max_results parameter limits the number of complete matches returned (not individual lines). |
| 198 | + When limited, the header shows "Found X matches (showing first Y of Z)". |
| 199 | +
|
| 200 | + Example usage: |
| 201 | + find_code_by_rule(yaml="id: x\\nlanguage: python\\nrule: {pattern: 'class $NAME'}", max_results=20) |
| 202 | + find_code_by_rule(yaml="...", output_format="json") # For full metadata |
| 203 | + """ |
| 204 | + if output_format not in ["text", "json"]: |
| 205 | + raise ValueError(f"Invalid output_format: {output_format}. Must be 'text' or 'json'.") |
| 206 | + |
| 207 | + args = ["--inline-rules", yaml] |
| 208 | + |
| 209 | + # Always get JSON internally for accurate match limiting |
| 210 | + result = run_ast_grep("scan", args + ["--json", project_folder]) |
| 211 | + matches = json.loads(result.stdout.strip() or "[]") |
| 212 | + |
| 213 | + # Apply max_results limit to complete matches |
| 214 | + total_matches = len(matches) |
| 215 | + if max_results is not None and total_matches > max_results: |
| 216 | + matches = matches[:max_results] |
209 | 217 |
|
210 |
| - # Apply max_results limit to complete matches |
211 |
| - total_matches = len(matches) |
212 |
| - if max_results is not None and total_matches > max_results: |
213 |
| - matches = matches[:max_results] |
| 218 | + if output_format == "text": |
| 219 | + if not matches: |
| 220 | + return "No matches found" |
| 221 | + text_output = format_matches_as_text(matches) |
| 222 | + header = f"Found {len(matches)} matches" |
| 223 | + if max_results is not None and total_matches > max_results: |
| 224 | + header += f" (showing first {max_results} of {total_matches})" |
| 225 | + return header + ":\n\n" + text_output |
| 226 | + return matches # type: ignore[no-any-return] |
214 | 227 |
|
215 |
| - if output_format == "text": |
216 |
| - if not matches: |
217 |
| - return "No matches found" |
218 |
| - text_output = format_matches_as_text(matches) |
219 |
| - header = f"Found {len(matches)} matches" |
220 |
| - if max_results is not None and total_matches > max_results: |
221 |
| - header += f" (showing first {max_results} of {total_matches})" |
222 |
| - return header + ":\n\n" + text_output |
223 |
| - return matches # type: ignore[no-any-return] |
224 | 228 |
|
225 | 229 | def format_matches_as_text(matches: List[dict]) -> str:
|
226 | 230 | """Convert JSON matches to LLM-friendly text format.
|
@@ -248,6 +252,29 @@ def format_matches_as_text(matches: List[dict]) -> str:
|
248 | 252 |
|
249 | 253 | return '\n\n'.join(output_blocks)
|
250 | 254 |
|
| 255 | +def get_supported_languages() -> List[str]: |
| 256 | + """Get all supported languages as a field description string.""" |
| 257 | + languages = [ # https://ast-grep.github.io/reference/languages.html |
| 258 | + "bash", "c", "cpp", "csharp", "css", "elixir", "go", "haskell", |
| 259 | + "html", "java", "javascript", "json", "jsx", "kotlin", "lua", |
| 260 | + "nix", "php", "python", "ruby", "rust", "scala", "solidity", |
| 261 | + "swift", "tsx", "typescript", "yaml" |
| 262 | + ] |
| 263 | + |
| 264 | + # Check for custom languages in config file |
| 265 | + # https://ast-grep.github.io/advanced/custom-language.html#register-language-in-sgconfig-yml |
| 266 | + if CONFIG_PATH and os.path.exists(CONFIG_PATH): |
| 267 | + try: |
| 268 | + with open(CONFIG_PATH, 'r') as f: |
| 269 | + config = yaml.safe_load(f) |
| 270 | + if config and 'customLanguages' in config: |
| 271 | + custom_langs = list(config['customLanguages'].keys()) |
| 272 | + languages += custom_langs |
| 273 | + except Exception: |
| 274 | + pass |
| 275 | + |
| 276 | + return sorted(set(languages)) |
| 277 | + |
251 | 278 | def run_command(args: List[str], input_text: Optional[str] = None) -> subprocess.CompletedProcess:
|
252 | 279 | try:
|
253 | 280 | # On Windows, if ast-grep is installed via npm, it's a batch file
|
@@ -281,7 +308,8 @@ def run_mcp_server() -> None:
|
281 | 308 | Run the MCP server.
|
282 | 309 | This function is used to start the MCP server when this script is run directly.
|
283 | 310 | """
|
284 |
| - parse_args_and_get_config() |
| 311 | + parse_args_and_get_config() # sets CONFIG_PATH |
| 312 | + register_mcp_tools() # tools defined *after* CONFIG_PATH is known |
285 | 313 | mcp.run(transport="stdio")
|
286 | 314 |
|
287 | 315 | if __name__ == "__main__":
|
|
0 commit comments