Skip to main content

project_map_cli_rust/mcp/
server.rs

1use tokio::io::{self, AsyncBufReadExt, BufReader, AsyncWriteExt};
2use serde::Deserialize;
3use serde_json::{json, Value};
4use crate::error::Result;
5use crate::core::query_engine::QueryEngine;
6use std::path::Path;
7
8#[derive(Deserialize, Debug)]
9struct JsonRpcRequest {
10    #[serde(rename = "jsonrpc")]
11    _jsonrpc: String,
12    id: Option<Value>,
13    method: String,
14    params: Option<Value>,
15}
16
17pub struct McpServer {
18    engine: Option<QueryEngine>,
19}
20
21impl McpServer {
22    pub fn new() -> Self {
23        let engine = QueryEngine::load(Path::new(".project-map/latest/.project-map.json")).ok();
24        Self { engine }
25    }
26
27    pub async fn run(&mut self) -> Result<()> {
28        let stdin = io::stdin();
29        let mut reader = BufReader::new(stdin).lines();
30        let mut stdout = io::stdout();
31
32        while let Some(line) = reader.next_line().await? {
33            let req: JsonRpcRequest = match serde_json::from_str(&line) {
34                Ok(r) => r,
35                Err(_) => continue,
36            };
37
38            let response = self.handle_request(req).await;
39            let response_json = serde_json::to_string(&response)?;
40            stdout.write_all(response_json.as_bytes()).await?;
41            stdout.write_all(b"\n").await?;
42            stdout.flush().await?;
43        }
44
45        Ok(())
46    }
47
48    async fn handle_request(&mut self, req: JsonRpcRequest) -> Value {
49        match req.method.as_str() {
50            "initialize" => json!({
51                "jsonrpc": "2.0",
52                "id": req.id,
53                "result": {
54                    "protocolVersion": "2024-11-05",
55                    "capabilities": {
56                        "tools": {}
57                    },
58                    "serverInfo": {
59                        "name": "project-map-cli-rust",
60                        "version": "0.1.0"
61                    }
62                }
63            }),
64            "notifications/initialized" => json!(null),
65            "tools/list" => json!({
66                "jsonrpc": "2.0",
67                "id": req.id,
68                "result": {
69                    "tools": [
70                        {
71                            "name": "pm_status",
72                            "description": "Returns current workspace context and available commands.",
73                            "inputSchema": { "type": "object", "properties": {} }
74                        },
75                        {
76                            "name": "pm_query",
77                            "description": "Search for symbols or get file context.",
78                            "inputSchema": {
79                                "type": "object",
80                                "properties": {
81                                    "query": { "type": "string" },
82                                    "path": { "type": "string" }
83                                }
84                            }
85                        },
86                        {
87                            "name": "pm_check_blast_radius",
88                            "description": "Identifies all components and files that depend on or import a specific symbol.",
89                            "inputSchema": {
90                                "type": "object",
91                                "properties": {
92                                    "path": { "type": "string" },
93                                    "symbol": { "type": "string" }
94                                },
95                                "required": ["path", "symbol"]
96                            }
97                        },
98                        {
99                            "name": "pm_plan",
100                            "description": "Analyze the architectural impact (fan-out) of a symbol before starting a refactor.",
101                            "inputSchema": {
102                                "type": "object",
103                                "properties": {
104                                    "symbol": { "type": "string" }
105                                },
106                                "required": ["symbol"]
107                            }
108                        },
109                        {
110                            "name": "pm_semantic_search",
111                            "description": "Search for logic using natural language keywords (e.g., 'auth', 'database').",
112                            "inputSchema": {
113                                "type": "object",
114                                "properties": {
115                                    "query": { "type": "string" }
116                                },
117                                "required": ["query"]
118                            }
119                        },
120                        {
121                            "name": "pm_fetch_symbol",
122                            "description": "Extract raw source code for a specific class or function.",
123                            "inputSchema": {
124                                "type": "object",
125                                "properties": {
126                                    "path": { "type": "string" },
127                                    "symbol": { "type": "string" }
128                                },
129                                "required": ["path", "symbol"]
130                            }
131                        },
132                        {
133                            "name": "pm_init",
134                            "description": "Refresh the map index after significant code changes to maintain discovery accuracy.",
135                            "inputSchema": { "type": "object", "properties": {} }
136                        }
137                    ]
138                }
139            }),
140            "tools/call" => self.handle_tool_call(req).await,
141            _ => json!({
142                "jsonrpc": "2.0",
143                "id": req.id,
144                "error": { "code": -32601, "message": "Method not found" }
145            }),
146        }
147    }
148
149    async fn handle_tool_call(&mut self, req: JsonRpcRequest) -> Value {
150        let params = req.params.as_ref().unwrap();
151        let tool_name = params.get("name").and_then(|v| v.as_str()).unwrap_or("");
152        let tool_args = params.get("arguments").cloned().unwrap_or(json!({}));
153
154        let text = match tool_name {
155            "pm_status" => {
156                if self.engine.is_some() {
157                    "Status: System healthy. Index is present.".to_string()
158                } else {
159                    "Status: Index missing. Run project-map build.".to_string()
160                }
161            }
162            "pm_query" => {
163                if let Some(ref engine) = self.engine {
164                    if let Some(q) = tool_args.get("query").and_then(|v| v.as_str()) {
165                        let matches = engine.find_symbols(q);
166                        format!("Matches: {}", matches.len())
167                    } else if let Some(p) = tool_args.get("path").and_then(|v| v.as_str()) {
168                        let symbols = engine.get_file_outline(p);
169                        format!("Symbols in {}: {}", p, symbols.len())
170                    } else {
171                        "Error: Provide query or path".to_string()
172                    }
173                } else {
174                    "Error: Index not loaded".to_string()
175                }
176            }
177            "pm_check_blast_radius" => {
178                if let Some(ref engine) = self.engine {
179                    let path = tool_args.get("path").and_then(|v| v.as_str()).unwrap_or("");
180                    let symbol = tool_args.get("symbol").and_then(|v| v.as_str()).unwrap_or("");
181                    let results = engine.check_blast_radius(path, symbol);
182
183                    if results.is_empty() {
184                        "No dependent components found.".to_string()
185                    } else {
186                        let mut unique_files = std::collections::HashSet::new();
187                        for r in &results { unique_files.insert(&r.path); }
188                        format!("Blast Radius for {}:\n- Total Impacted Nodes: {}\n- Unique Files: {}\n(Top 5: {})", 
189                            symbol, results.len(), unique_files.len(),
190                            results.iter().take(5).map(|r| r.name.as_str()).collect::<Vec<_>>().join(", "))
191                    }
192                } else {
193                    "Error: Index not loaded".to_string()
194                }
195            }
196            "pm_plan" => {
197                if let Some(ref engine) = self.engine {
198                    let symbol = tool_args.get("symbol").and_then(|v| v.as_str()).unwrap_or("");
199                    let impact = engine.analyze_impact(symbol);
200                    let blast = engine.check_blast_radius("", symbol);
201
202                    let mut unique_blast = std::collections::HashSet::new();
203                    for r in &blast { unique_blast.insert(&r.path); }
204
205                    format!("Architectural Plan for {}:\n- Fan-out (Dependencies): {} nodes\n- Fan-in (Dependents): {} nodes across {} files.", 
206                        symbol, impact.len(), blast.len(), unique_blast.len())
207                } else {
208                    "Error: Index not loaded".to_string()
209                }
210            }
211            "pm_semantic_search" => {
212                if let Some(ref engine) = self.engine {
213                    let query = tool_args.get("query").and_then(|v| v.as_str()).unwrap_or("");
214                    let matches = engine.find_symbols(query);
215                    let mut result = format!("Semantic Search Results ({}):", matches.len());
216                    for m in matches.iter().take(15) {
217                        result.push_str(&format!("\n- {}: {}", m.path, m.name));
218                    }
219                    result
220                } else {
221                    "Error: Index not loaded".to_string()
222                }
223            }
224            "pm_fetch_symbol" => {
225                if let Some(ref engine) = self.engine {
226                    let path = tool_args.get("path").and_then(|v| v.as_str()).unwrap_or("");
227                    let symbol = tool_args.get("symbol").and_then(|v| v.as_str()).unwrap_or("");
228                    if let Some(node) = engine.find_symbol_in_path(path, symbol) {
229                        if let Ok(content) = std::fs::read_to_string(&node.path) {
230                            let bytes = content.as_bytes();
231                            if node.start_byte < bytes.len() && node.end_byte <= bytes.len() {
232                                String::from_utf8_lossy(&bytes[node.start_byte..node.end_byte]).to_string()
233                            } else {
234                                "Error: Byte range out of bounds".to_string()
235                            }
236                        } else {
237                            "Error: Could not read file".to_string()
238                        }
239                    } else {
240                        "Error: Symbol not found".to_string()
241                    }
242                } else {
243                    "Error: Index not loaded".to_string()
244                }
245            }
246            "pm_init" => {
247                use crate::core::orchestrator::Orchestrator;
248                let mut orch = Orchestrator::new();
249                if orch.build_index(Path::new(".")).is_ok() && orch.save_index_versioned(Path::new(".project-map")).is_ok() {
250                    self.engine = QueryEngine::load(Path::new(".project-map/latest/.project-map.json")).ok();
251                    "Index refreshed successfully.".to_string()
252                } else {
253                    "Failed to refresh index.".to_string()
254                }
255            }
256
257            _ => "Error: Unknown tool".to_string(),
258        };
259
260        json!({
261            "jsonrpc": "2.0",
262            "id": req.id,
263            "result": {
264                "content": [
265                    { "type": "text", "text": text }
266                ]
267            }
268        })
269    }
270}