| 
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