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.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(&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                }
111            }),
112            "tools/call" => self.handle_tool_call(req).await,
113            _ => json!({
114                "jsonrpc": "2.0",
115                "id": req.id,
116                "error": { "code": -32601, "message": "Method not found" }
117            }),
118        }
119    }
120
121    async fn handle_tool_call(&self, req: JsonRpcRequest) -> Value {
122        let params = req.params.as_ref().unwrap();
123        let tool_name = params.get("name").and_then(|v| v.as_str()).unwrap_or("");
124        let tool_args = params.get("arguments").cloned().unwrap_or(json!({}));
125
126        let text = match tool_name {
127            "pm_status" => {
128                if self.engine.is_some() {
129                    "Status: System healthy. Index is present.".to_string()
130                } else {
131                    "Status: Index missing. Run project-map build.".to_string()
132                }
133            }
134            "pm_query" => {
135                if let Some(ref engine) = self.engine {
136                    if let Some(q) = tool_args.get("query").and_then(|v| v.as_str()) {
137                        let matches = engine.find_symbols(q);
138                        format!("Matches: {}", matches.len())
139                    } else if let Some(p) = tool_args.get("path").and_then(|v| v.as_str()) {
140                        let symbols = engine.get_file_outline(p);
141                        format!("Symbols in {}: {}", p, symbols.len())
142                    } else {
143                        "Error: Provide query or path".to_string()
144                    }
145                } else {
146                    "Error: Index not loaded".to_string()
147                }
148            }
149            "pm_check_blast_radius" => {
150                if let Some(ref engine) = self.engine {
151                    let path = tool_args.get("path").and_then(|v| v.as_str()).unwrap_or("");
152                    let symbol = tool_args.get("symbol").and_then(|v| v.as_str()).unwrap_or("");
153                    let results = engine.check_blast_radius(path, symbol);
154
155                    if results.is_empty() {
156                        "No dependent components found.".to_string()
157                    } else {
158                        let mut unique_files = std::collections::HashSet::new();
159                        for r in &results { unique_files.insert(&r.path); }
160                        format!("Blast Radius for {}:\n- Total Impacted Nodes: {}\n- Unique Files: {}\n(Top 5: {})", 
161                            symbol, results.len(), unique_files.len(),
162                            results.iter().take(5).map(|r| r.name.as_str()).collect::<Vec<_>>().join(", "))
163                    }
164                } else {
165                    "Error: Index not loaded".to_string()
166                }
167            }
168            "pm_plan" => {
169                if let Some(ref engine) = self.engine {
170                    let symbol = tool_args.get("symbol").and_then(|v| v.as_str()).unwrap_or("");
171                    let impact = engine.analyze_impact(symbol);
172                    let blast = engine.check_blast_radius("", symbol);
173
174                    let mut unique_blast = std::collections::HashSet::new();
175                    for r in &blast { unique_blast.insert(&r.path); }
176
177                    format!("Architectural Plan for {}:\n- Fan-out (Dependencies): {} nodes\n- Fan-in (Dependents): {} nodes across {} files.", 
178                        symbol, impact.len(), blast.len(), unique_blast.len())
179                } else {
180                    "Error: Index not loaded".to_string()
181                }
182            }
183
184            _ => "Error: Unknown tool".to_string(),
185        };
186
187        json!({
188            "jsonrpc": "2.0",
189            "id": req.id,
190            "result": {
191                "content": [
192                    { "type": "text", "text": text }
193                ]
194            }
195        })
196    }
197}