1
- from typing import Any , List , Optional
2
- from mcp .server .fastmcp import FastMCP
3
- import subprocess
4
- from pydantic import Field
5
- import json
6
- from enum import Enum
7
1
import argparse
2
+ import json
8
3
import os
4
+ import subprocess
9
5
import sys
6
+ from enum import Enum
7
+ from typing import Any , List , Optional
8
+
9
+ from mcp .server .fastmcp import FastMCP
10
+ from pydantic import Field
10
11
11
- # Determine how the script was invoked
12
- if sys .argv [0 ].endswith ('main.py' ):
13
- # Direct execution: python main.py
14
- prog = 'python main.py'
15
- else :
16
- # Installed script execution (via uvx, pip install, etc.)
17
- prog = None # Let argparse use the default
18
-
19
- # Parse command-line arguments
20
- parser = argparse .ArgumentParser (
21
- prog = prog ,
22
- description = 'ast-grep MCP Server - Provides structural code search capabilities via Model Context Protocol' ,
23
- epilog = '''
12
+ # Global variable for config path (will be set by parse_args_and_get_config)
13
+ CONFIG_PATH = None
14
+
15
+ def parse_args_and_get_config ():
16
+ """Parse command-line arguments and determine config path."""
17
+ global CONFIG_PATH
18
+
19
+ # Determine how the script was invoked
20
+ prog = None
21
+ if sys .argv [0 ].endswith ('main.py' ):
22
+ # Direct execution: python main.py
23
+ prog = 'python main.py'
24
+
25
+ # Parse command-line arguments
26
+ parser = argparse .ArgumentParser (
27
+ prog = prog ,
28
+ description = 'ast-grep MCP Server - Provides structural code search capabilities via Model Context Protocol' ,
29
+ epilog = '''
24
30
environment variables:
25
31
AST_GREP_CONFIG Path to sgconfig.yaml file (overridden by --config flag)
26
32
27
33
For more information, see: https://github.com/ast-grep/ast-grep-mcp
28
- ''' ,
29
- formatter_class = argparse .RawDescriptionHelpFormatter
30
- )
31
- parser .add_argument (
32
- '--config' ,
33
- type = str ,
34
- metavar = 'PATH' ,
35
- help = 'Path to sgconfig.yaml file for customizing ast-grep behavior (language mappings, rule directories, etc.)'
36
- )
37
- args = parser .parse_args ()
38
-
39
- # Determine config path with precedence: --config flag > AST_GREP_CONFIG env > None
40
- CONFIG_PATH = None
41
- if args .config :
42
- if not os .path .exists (args .config ):
43
- print (f"Error: Config file '{ args .config } ' does not exist" )
44
- sys .exit (1 )
45
- CONFIG_PATH = args .config
46
- elif os .environ .get ('AST_GREP_CONFIG' ):
47
- env_config = os .environ .get ('AST_GREP_CONFIG' )
48
- if not os .path .exists (env_config ):
49
- print (f"Error: Config file '{ env_config } ' specified in AST_GREP_CONFIG does not exist" )
50
- sys .exit (1 )
51
- CONFIG_PATH = env_config
34
+ ''' ,
35
+ formatter_class = argparse .RawDescriptionHelpFormatter
36
+ )
37
+ parser .add_argument (
38
+ '--config' ,
39
+ type = str ,
40
+ metavar = 'PATH' ,
41
+ help = 'Path to sgconfig.yaml file for customizing ast-grep behavior (language mappings, rule directories, etc.)'
42
+ )
43
+ args = parser .parse_args ()
44
+
45
+ # Determine config path with precedence: --config flag > AST_GREP_CONFIG env > None
46
+ if args .config :
47
+ if not os .path .exists (args .config ):
48
+ print (f"Error: Config file '{ args .config } ' does not exist" )
49
+ sys .exit (1 )
50
+ CONFIG_PATH = args .config
51
+ elif os .environ .get ('AST_GREP_CONFIG' ):
52
+ env_config = os .environ .get ('AST_GREP_CONFIG' )
53
+ if env_config and not os .path .exists (env_config ):
54
+ print (f"Error: Config file '{ env_config } ' specified in AST_GREP_CONFIG does not exist" )
55
+ sys .exit (1 )
56
+ CONFIG_PATH = env_config
52
57
53
58
# Initialize FastMCP server
54
59
mcp = FastMCP ("ast-grep" )
@@ -62,7 +67,7 @@ class DumpFormat(Enum):
62
67
def dump_syntax_tree (
63
68
code : str = Field (description = "The code you need" ),
64
69
language : str = Field (description = "The language of the code" ),
65
- format : DumpFormat = Field (description = "Code dump format. Available values: pattern, ast, cst" , default = "cst" ),
70
+ format : DumpFormat = Field (description = "Code dump format. Available values: pattern, ast, cst" , default = DumpFormat . CST ),
66
71
) -> str :
67
72
"""
68
73
Dump code's syntax structure or dump a query's pattern structure.
@@ -75,7 +80,7 @@ def dump_syntax_tree(
75
80
Internally calls: ast-grep run --pattern <code> --lang <language> --debug-query=<format>
76
81
"""
77
82
result = run_ast_grep ("run" , ["--pattern" , code , "--lang" , language , f"--debug-query={ format .value } " ])
78
- return result .stderr .strip ()
83
+ return result .stderr .strip () # type: ignore[no-any-return]
79
84
80
85
@mcp .tool ()
81
86
def test_match_code_rule (
@@ -92,7 +97,7 @@ def test_match_code_rule(
92
97
matches = json .loads (result .stdout .strip ())
93
98
if not matches :
94
99
raise ValueError ("No matches found for the given code and rule. Try adding `stopBy: end` to your inside/has rule." )
95
- return matches
100
+ return matches # type: ignore[no-any-return]
96
101
97
102
@mcp .tool ()
98
103
def find_code (
@@ -131,7 +136,7 @@ def find_code(
131
136
# Limit results if max_results is specified
132
137
if max_results is not None and len (matches ) > max_results :
133
138
matches = matches [:max_results ]
134
- return matches
139
+ return matches # type: ignore[no-any-return]
135
140
else :
136
141
# Text format - return plain text output
137
142
result = run_ast_grep ("run" , args + [project_folder ])
@@ -150,7 +155,7 @@ def find_code(
150
155
else :
151
156
header = f"Found { len (non_empty_lines )} matches:\n "
152
157
output = header + output
153
- return output
158
+ return output # type: ignore[no-any-return]
154
159
155
160
@mcp .tool ()
156
161
def find_code_by_rule (
@@ -188,7 +193,7 @@ def find_code_by_rule(
188
193
# Limit results if max_results is specified
189
194
if max_results is not None and len (matches ) > max_results :
190
195
matches = matches [:max_results ]
191
- return matches
196
+ return matches # type: ignore[no-any-return]
192
197
else :
193
198
# Text format - return plain text output
194
199
result = run_ast_grep ("scan" , args + [project_folder ])
@@ -207,16 +212,21 @@ def find_code_by_rule(
207
212
else :
208
213
header = f"Found { len (non_empty_lines )} matches:\n "
209
214
output = header + output
210
- return output
215
+ return output # type: ignore[no-any-return]
211
216
212
217
def run_command (args : List [str ], input_text : Optional [str ] = None ) -> subprocess .CompletedProcess :
213
218
try :
219
+ # On Windows, if ast-grep is installed via npm, it's a batch file
220
+ # that requires shell=True to execute properly
221
+ use_shell = (sys .platform == "win32" and args [0 ] == "ast-grep" )
222
+
214
223
result = subprocess .run (
215
224
args ,
216
225
capture_output = True ,
217
226
input = input_text ,
218
227
text = True ,
219
- check = True # Raises CalledProcessError if return code is non-zero
228
+ check = True , # Raises CalledProcessError if return code is non-zero
229
+ shell = use_shell
220
230
)
221
231
return result
222
232
except subprocess .CalledProcessError as e :
@@ -237,6 +247,7 @@ def run_mcp_server() -> None:
237
247
Run the MCP server.
238
248
This function is used to start the MCP server when this script is run directly.
239
249
"""
250
+ parse_args_and_get_config ()
240
251
mcp .run (transport = "stdio" )
241
252
242
253
if __name__ == "__main__" :
0 commit comments