Skip to main content

code_baseline/
mcp.rs

1use crate::cli::toml_config::TomlConfig;
2use crate::presets;
3use crate::scan;
4use serde_json::json;
5use std::io::{self, BufRead, Write};
6use std::path::{Path, PathBuf};
7
8/// Run a simple MCP-compatible server over stdio.
9///
10/// Reads JSON-RPC requests from stdin, processes them, and writes
11/// JSON-RPC responses to stdout. Supports the MCP protocol for
12/// tool discovery and execution.
13pub fn run_mcp_server(config_path: &Path) {
14    let stdin = io::stdin();
15    let mut stdout = io::stdout();
16
17    // Read line-delimited JSON-RPC messages
18    for line in stdin.lock().lines() {
19        let line = match line {
20            Ok(l) => l,
21            Err(_) => break,
22        };
23
24        if line.trim().is_empty() {
25            continue;
26        }
27
28        let request: serde_json::Value = match serde_json::from_str(&line) {
29            Ok(v) => v,
30            Err(e) => {
31                let error_response = json!({
32                    "jsonrpc": "2.0",
33                    "id": null,
34                    "error": { "code": -32700, "message": format!("Parse error: {}", e) }
35                });
36                let _ = writeln!(stdout, "{}", error_response);
37                let _ = stdout.flush();
38                continue;
39            }
40        };
41
42        let id = request.get("id").cloned();
43        let method = request.get("method").and_then(|m| m.as_str()).unwrap_or("");
44        let params = request.get("params").cloned().unwrap_or(json!({}));
45
46        let response = match method {
47            "initialize" => handle_initialize(id.clone()),
48            "tools/list" => handle_tools_list(id.clone()),
49            "tools/call" => handle_tools_call(id.clone(), &params, config_path),
50            "notifications/initialized" | "notifications/cancelled" => continue,
51            _ => json!({
52                "jsonrpc": "2.0",
53                "id": id,
54                "error": { "code": -32601, "message": format!("Unknown method: {}", method) }
55            }),
56        };
57
58        let _ = writeln!(stdout, "{}", response);
59        let _ = stdout.flush();
60    }
61}
62
63fn handle_initialize(id: Option<serde_json::Value>) -> serde_json::Value {
64    json!({
65        "jsonrpc": "2.0",
66        "id": id,
67        "result": {
68            "protocolVersion": "2024-11-05",
69            "capabilities": {
70                "tools": {}
71            },
72            "serverInfo": {
73                "name": "baseline",
74                "version": env!("CARGO_PKG_VERSION")
75            }
76        }
77    })
78}
79
80fn handle_tools_list(id: Option<serde_json::Value>) -> serde_json::Value {
81    json!({
82        "jsonrpc": "2.0",
83        "id": id,
84        "result": {
85            "tools": [
86                {
87                    "name": "baseline_scan",
88                    "description": "Scan files for rule violations. Returns structured violations with fix suggestions.",
89                    "inputSchema": {
90                        "type": "object",
91                        "properties": {
92                            "paths": {
93                                "type": "array",
94                                "items": { "type": "string" },
95                                "description": "File or directory paths to scan"
96                            },
97                            "content": {
98                                "type": "string",
99                                "description": "Inline file content to scan (alternative to paths)"
100                            },
101                            "filename": {
102                                "type": "string",
103                                "description": "Virtual filename for glob matching when using content"
104                            }
105                        }
106                    }
107                },
108                {
109                    "name": "baseline_list_rules",
110                    "description": "List all configured rules and their descriptions.",
111                    "inputSchema": {
112                        "type": "object",
113                        "properties": {}
114                    }
115                }
116            ]
117        }
118    })
119}
120
121fn handle_tools_call(
122    id: Option<serde_json::Value>,
123    params: &serde_json::Value,
124    config_path: &Path,
125) -> serde_json::Value {
126    let tool_name = params
127        .get("name")
128        .and_then(|n| n.as_str())
129        .unwrap_or("");
130
131    let arguments = params.get("arguments").cloned().unwrap_or(json!({}));
132
133    match tool_name {
134        "baseline_scan" => handle_scan(&id, &arguments, config_path),
135        "baseline_list_rules" => handle_list_rules(&id, config_path),
136        _ => json!({
137            "jsonrpc": "2.0",
138            "id": id,
139            "error": { "code": -32602, "message": format!("Unknown tool: {}", tool_name) }
140        }),
141    }
142}
143
144fn handle_scan(
145    id: &Option<serde_json::Value>,
146    arguments: &serde_json::Value,
147    config_path: &Path,
148) -> serde_json::Value {
149    // Check for inline content mode
150    if let Some(content) = arguments.get("content").and_then(|c| c.as_str()) {
151        let filename = arguments
152            .get("filename")
153            .and_then(|f| f.as_str())
154            .unwrap_or("stdin.tsx");
155
156        match scan::run_scan_stdin(config_path, content, filename) {
157            Ok(result) => {
158                let violations = format_violations_json(&result);
159                json!({
160                    "jsonrpc": "2.0",
161                    "id": id,
162                    "result": {
163                        "content": [{ "type": "text", "text": violations.to_string() }]
164                    }
165                })
166            }
167            Err(e) => json!({
168                "jsonrpc": "2.0",
169                "id": id,
170                "result": {
171                    "content": [{ "type": "text", "text": format!("Error: {}", e) }],
172                    "isError": true
173                }
174            }),
175        }
176    } else {
177        // File paths mode
178        let paths: Vec<PathBuf> = arguments
179            .get("paths")
180            .and_then(|p| p.as_array())
181            .map(|arr| {
182                arr.iter()
183                    .filter_map(|v| v.as_str().map(PathBuf::from))
184                    .collect()
185            })
186            .unwrap_or_else(|| vec![PathBuf::from(".")]);
187
188        match scan::run_scan(config_path, &paths) {
189            Ok(result) => {
190                let violations = format_violations_json(&result);
191                json!({
192                    "jsonrpc": "2.0",
193                    "id": id,
194                    "result": {
195                        "content": [{ "type": "text", "text": violations.to_string() }]
196                    }
197                })
198            }
199            Err(e) => json!({
200                "jsonrpc": "2.0",
201                "id": id,
202                "result": {
203                    "content": [{ "type": "text", "text": format!("Error: {}", e) }],
204                    "isError": true
205                }
206            }),
207        }
208    }
209}
210
211fn handle_list_rules(
212    id: &Option<serde_json::Value>,
213    config_path: &Path,
214) -> serde_json::Value {
215    let config_text = match std::fs::read_to_string(config_path) {
216        Ok(c) => c,
217        Err(e) => {
218            return json!({
219                "jsonrpc": "2.0",
220                "id": id,
221                "result": {
222                    "content": [{ "type": "text", "text": format!("Error reading config: {}", e) }],
223                    "isError": true
224                }
225            });
226        }
227    };
228
229    let toml_config: TomlConfig = match toml::from_str(&config_text) {
230        Ok(c) => c,
231        Err(e) => {
232            return json!({
233                "jsonrpc": "2.0",
234                "id": id,
235                "result": {
236                    "content": [{ "type": "text", "text": format!("Error parsing config: {}", e) }],
237                    "isError": true
238                }
239            });
240        }
241    };
242
243    let mut resolved = match presets::resolve_rules(&toml_config.baseline.extends, &toml_config.rule) {
244        Ok(r) => r,
245        Err(e) => {
246            return json!({
247                "jsonrpc": "2.0",
248                "id": id,
249                "result": {
250                    "content": [{ "type": "text", "text": format!("Error resolving rules: {}", e) }],
251                    "isError": true
252                }
253            });
254        }
255    };
256
257    match presets::resolve_scoped_rules(&toml_config.baseline.scoped, &toml_config.rule) {
258        Ok(scoped) => resolved.extend(scoped),
259        Err(e) => {
260            return json!({
261                "jsonrpc": "2.0",
262                "id": id,
263                "result": {
264                    "content": [{ "type": "text", "text": format!("Error resolving scoped rules: {}", e) }],
265                    "isError": true
266                }
267            });
268        }
269    }
270
271    let rules: Vec<serde_json::Value> = resolved
272        .iter()
273        .map(|r| {
274            json!({
275                "id": r.id,
276                "type": r.rule_type,
277                "severity": r.severity,
278                "glob": r.glob,
279                "message": r.message,
280            })
281        })
282        .collect();
283
284    let text = serde_json::to_string_pretty(&json!({ "rules": rules })).unwrap();
285
286    json!({
287        "jsonrpc": "2.0",
288        "id": id,
289        "result": {
290            "content": [{ "type": "text", "text": text }]
291        }
292    })
293}
294
295fn format_violations_json(result: &scan::ScanResult) -> serde_json::Value {
296    use crate::config::Severity;
297
298    let violations: Vec<serde_json::Value> = result
299        .violations
300        .iter()
301        .map(|v| {
302            let mut obj = json!({
303                "rule_id": v.rule_id,
304                "severity": match v.severity {
305                    Severity::Error => "error",
306                    Severity::Warning => "warning",
307                },
308                "file": v.file.display().to_string(),
309                "line": v.line,
310                "column": v.column,
311                "message": v.message,
312                "suggest": v.suggest,
313            });
314
315            if let Some(ref fix) = v.fix {
316                obj["fix"] = json!({ "old": fix.old, "new": fix.new });
317            }
318
319            obj
320        })
321        .collect();
322
323    json!({
324        "violations": violations,
325        "summary": {
326            "total": result.violations.len(),
327            "errors": result.violations.iter().filter(|v| v.severity == Severity::Error).count(),
328            "warnings": result.violations.iter().filter(|v| v.severity == Severity::Warning).count(),
329            "files_scanned": result.files_scanned,
330        }
331    })
332}
333
334#[cfg(test)]
335mod tests {
336    use super::*;
337    use crate::config::Severity;
338    use crate::rules::Violation;
339    use std::collections::HashMap;
340    use std::path::PathBuf;
341
342    #[test]
343    fn initialize_returns_protocol_version() {
344        let resp = handle_initialize(Some(json!(1)));
345        assert_eq!(resp["jsonrpc"], "2.0");
346        assert_eq!(resp["id"], 1);
347        assert_eq!(resp["result"]["protocolVersion"], "2024-11-05");
348        assert_eq!(resp["result"]["serverInfo"]["name"], "baseline");
349    }
350
351    #[test]
352    fn tools_list_returns_both_tools() {
353        let resp = handle_tools_list(Some(json!(2)));
354        assert_eq!(resp["jsonrpc"], "2.0");
355        let tools = resp["result"]["tools"].as_array().unwrap();
356        assert_eq!(tools.len(), 2);
357        assert_eq!(tools[0]["name"], "baseline_scan");
358        assert_eq!(tools[1]["name"], "baseline_list_rules");
359    }
360
361    #[test]
362    fn format_violations_empty() {
363        let result = scan::ScanResult {
364            violations: vec![],
365            files_scanned: 3,
366            rules_loaded: 2,
367            ratchet_counts: HashMap::new(),
368            changed_files_count: None,
369            base_ref: None,
370        };
371        let json = format_violations_json(&result);
372        assert_eq!(json["summary"]["total"], 0);
373        assert_eq!(json["summary"]["files_scanned"], 3);
374        assert!(json["violations"].as_array().unwrap().is_empty());
375    }
376
377    #[test]
378    fn format_violations_with_fix() {
379        let result = scan::ScanResult {
380            violations: vec![Violation {
381                rule_id: "test-rule".into(),
382                severity: Severity::Error,
383                file: PathBuf::from("test.tsx"),
384                line: Some(5),
385                column: Some(10),
386                message: "bad class".into(),
387                suggest: Some("use good class".into()),
388                source_line: None,
389                fix: Some(crate::rules::Fix {
390                    old: "bg-red-500".into(),
391                    new: "bg-destructive".into(),
392                }),
393            }],
394            files_scanned: 1,
395            rules_loaded: 1,
396            ratchet_counts: HashMap::new(),
397            changed_files_count: None,
398            base_ref: None,
399        };
400        let json = format_violations_json(&result);
401        assert_eq!(json["summary"]["total"], 1);
402        assert_eq!(json["summary"]["errors"], 1);
403        let v = &json["violations"][0];
404        assert_eq!(v["rule_id"], "test-rule");
405        assert_eq!(v["fix"]["old"], "bg-red-500");
406        assert_eq!(v["fix"]["new"], "bg-destructive");
407    }
408
409    #[test]
410    fn format_violations_counts_severities() {
411        let result = scan::ScanResult {
412            violations: vec![
413                Violation {
414                    rule_id: "r1".into(),
415                    severity: Severity::Error,
416                    file: PathBuf::from("a.ts"),
417                    line: Some(1),
418                    column: None,
419                    message: "err".into(),
420                    suggest: None,
421                    source_line: None,
422                    fix: None,
423                },
424                Violation {
425                    rule_id: "r2".into(),
426                    severity: Severity::Warning,
427                    file: PathBuf::from("b.ts"),
428                    line: Some(2),
429                    column: None,
430                    message: "warn".into(),
431                    suggest: None,
432                    source_line: None,
433                    fix: None,
434                },
435            ],
436            files_scanned: 2,
437            rules_loaded: 2,
438            ratchet_counts: HashMap::new(),
439            changed_files_count: None,
440            base_ref: None,
441        };
442        let json = format_violations_json(&result);
443        assert_eq!(json["summary"]["errors"], 1);
444        assert_eq!(json["summary"]["warnings"], 1);
445        assert_eq!(json["summary"]["total"], 2);
446    }
447
448    #[test]
449    fn unknown_tool_returns_error() {
450        let resp = handle_tools_call(
451            Some(json!(3)),
452            &json!({ "name": "nonexistent_tool", "arguments": {} }),
453            std::path::Path::new("baseline.toml"),
454        );
455        assert!(resp["error"].is_object());
456        assert_eq!(resp["error"]["code"], -32602);
457    }
458}