From 147f53f611cce2c1ef293f5a5cff65c80cc36546 Mon Sep 17 00:00:00 2001 From: Rakshat28 Date: Sun, 24 May 2026 03:28:19 +0530 Subject: [PATCH] MCP --- Cargo.lock | 1 + Cargo.toml | 1 + src/main.rs | 42 ++++- src/mcp.rs | 428 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/memory.rs | 20 +++ 5 files changed, 491 insertions(+), 1 deletion(-) create mode 100644 src/mcp.rs diff --git a/Cargo.lock b/Cargo.lock index 6fafc47..8df50b6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -383,6 +383,7 @@ dependencies = [ "rusqlite", "scopeguard", "serde", + "serde_json", "similar", "tempfile", "thiserror", diff --git a/Cargo.toml b/Cargo.toml index 341fc2c..8ab3f01 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,7 @@ memmap2 = "0.9" bincode = { version = "2", features = ["serde"] } serde = { version = "1", features = ["derive"] } rusqlite = { version = "0.31", features = ["bundled"] } +serde_json = "1" ratatui = "0.27" crossterm = { version = "0.27", features = ["event-stream"] } diff --git a/src/main.rs b/src/main.rs index 26e1810..7f22979 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,6 +16,7 @@ use rayon::prelude::*; mod bloom; mod index; +mod mcp; #[allow(dead_code)] mod memory; mod output; @@ -29,6 +30,7 @@ pub mod walker; use bloom::BloomFilter; use index::{index_path_for_root, load_index, save_index, IndexEntry, IndexManifest}; +use mcp::{run as run_mcp, McpConfig}; use memory::{memory_db_path, FileRow, MemoryDb, SymbolKind, SymbolRow}; use output::{print_lookup_results, print_match, print_summary, resolve_color_mode, ColorMode}; use parser::{detect_language, get_all_languages, parse_file_with_metadata}; @@ -91,6 +93,9 @@ struct App { #[command(flatten)] cli: Cli, + #[arg(long = "mcp", default_value_t = false, help = "Run the MCP server over stdio")] + mcp: bool, + #[command(subcommand)] command: Option, } @@ -250,11 +255,19 @@ struct SearchOutcome { impl Cli { fn validate(&self) -> std::result::Result<(), String> { + self.validate_internal(false) + } + + fn validate_allowing_empty_query(&self) -> std::result::Result<(), String> { + self.validate_internal(true) + } + + fn validate_internal(&self, allow_empty_query: bool) -> std::result::Result<(), String> { if self.generate_completions.is_some() { return Ok(()); } - if self.query.iter().all(|q| q.trim().is_empty()) { + if !allow_empty_query && self.query.iter().all(|q| q.trim().is_empty()) { return Err("at least one query string must not be empty".to_string()); } @@ -486,6 +499,13 @@ fn run_lookup_mode(args: &LookupArgs) { print_lookup_results(&results, &color, &mut stdout); } +fn resolve_root_path(path: &Path) -> PathBuf { + match fs::canonicalize(path) { + Ok(resolved) => resolved, + Err(_) => path.to_path_buf(), + } +} + fn resolve_lang(lang_str: &str) -> Language { match lang_str { "rust" => Language::Rust, @@ -850,6 +870,26 @@ fn main() { return; } + if app.mcp { + if let Err(message) = cli.validate_allowing_empty_query() { + eprintln!("error: {message}"); + process::exit(1); + } + + let root_path = resolve_root_path(&cli.path); + let mcp_config = McpConfig { + root_path: root_path.clone(), + db_path: memory_db_path(&root_path), + lang_mode: resolve_lang_mode(&cli.lang), + }; + + if let Err(error) = run_mcp(mcp_config) { + eprintln!("error: {error}"); + process::exit(1); + } + return; + } + if let Some(shell) = cli.generate_completions { let mut cmd = App::command(); generate(shell, &mut cmd, "dora", &mut std::io::stdout()); diff --git a/src/mcp.rs b/src/mcp.rs new file mode 100644 index 0000000..888e4d3 --- /dev/null +++ b/src/mcp.rs @@ -0,0 +1,428 @@ +use crate::index::{index_path_for_root, load_index, IndexManifest}; +use crate::memory::{FileRow, MemoryDb, SymbolRow}; +use crate::output::ColorMode; +use crate::parser::get_all_languages; +use crate::query; +use crate::sieve::build_query_trigram_set; +use crate::types::{AppError, LangMode, Language, MatchResult, Result, SearchConfig}; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::io::{BufRead, Write}; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +#[derive(Clone)] +pub struct McpConfig { + pub root_path: PathBuf, + pub db_path: PathBuf, + pub lang_mode: LangMode, +} + +#[derive(Debug, Deserialize)] +struct RpcRequest { + jsonrpc: String, + #[serde(default)] + id: Option, + method: String, + #[serde(default)] + params: Option, +} + +#[derive(Debug, Serialize)] +struct RpcResponse { + jsonrpc: &'static str, + id: Value, + #[serde(skip_serializing_if = "Option::is_none")] + result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +#[derive(Debug, Serialize)] +struct RpcError { + code: i32, + message: String, + #[serde(skip_serializing_if = "Option::is_none")] + data: Option, +} + +#[derive(Debug, Deserialize)] +struct SearchAstParams { + query: String, +} + +#[derive(Debug, Deserialize)] +struct LookupSymbolParams { + name: String, +} + +#[derive(Debug, Serialize)] +struct SearchAstItem { + file_path: String, + capture_name: String, + matched_text: String, + start_line: usize, + start_col: usize, + end_line: usize, + end_col: usize, +} + +#[derive(Debug, Serialize)] +struct LookupSymbolItem { + file_path: String, + kind: String, + name: String, + start_line: usize, + start_col: usize, + end_line: usize, + end_col: usize, + signature: Option, +} + +pub fn run(config: McpConfig) -> Result<()> { + let stdin = std::io::stdin(); + let stdout = std::io::stdout(); + let mut input = stdin.lock(); + let mut output = stdout.lock(); + let mut line = String::new(); + + loop { + line.clear(); + let read = input.read_line(&mut line).map_err(AppError::IoError)?; + if read == 0 { + break; + } + if line.trim().is_empty() { + continue; + } + + let response = handle_line(&config, line.trim_end()); + output.write_all(response.as_bytes()).map_err(AppError::IoError)?; + output.write_all(b"\n").map_err(AppError::IoError)?; + output.flush().map_err(AppError::IoError)?; + } + + Ok(()) +} + +fn handle_line(config: &McpConfig, line: &str) -> String { + match serde_json::from_str::(line) { + Ok(request) => response_to_string(handle_request(config, request)), + Err(error) => response_to_string(RpcResponse { + jsonrpc: "2.0", + id: Value::Null, + result: None, + error: Some(RpcError { + code: -32700, + message: "parse error".to_string(), + data: Some(Value::String(error.to_string())), + }), + }), + } +} + +fn handle_request(config: &McpConfig, request: RpcRequest) -> RpcResponse { + let id = request.id.unwrap_or(Value::Null); + + if request.jsonrpc != "2.0" { + return rpc_error( + id, + -32600, + "invalid request", + Some(Value::String("jsonrpc must be \"2.0\"".to_string())), + ); + } + + match request.method.as_str() { + "search_ast" => match parse_params::(request.params) { + Ok(params) if params.query.trim().is_empty() => rpc_error( + id, + -32602, + "invalid params", + Some(Value::String("query must not be empty".to_string())), + ), + Ok(params) => match search_ast(config, ¶ms.query) { + Ok(results) => rpc_ok(id, json!(results)), + Err(error) => rpc_app_error(id, error), + }, + Err(message) => rpc_error(id, -32602, "invalid params", Some(Value::String(message))), + }, + "lookup_symbol" => match parse_params::(request.params) { + Ok(params) if params.name.trim().is_empty() => rpc_error( + id, + -32602, + "invalid params", + Some(Value::String("name must not be empty".to_string())), + ), + Ok(params) => match lookup_symbol(config, ¶ms.name) { + Ok(results) => rpc_ok(id, json!(results)), + Err(error) => rpc_app_error(id, error), + }, + Err(message) => rpc_error(id, -32602, "invalid params", Some(Value::String(message))), + }, + _ => rpc_error(id, -32601, "method not found", Some(Value::String(request.method))), + } +} + +fn parse_params(params: Option) -> std::result::Result +where + T: for<'de> Deserialize<'de>, +{ + let value = params.ok_or_else(|| "missing params".to_string())?; + serde_json::from_value(value).map_err(|error| error.to_string()) +} + +fn rpc_ok(id: Value, result: Value) -> RpcResponse { + RpcResponse { jsonrpc: "2.0", id, result: Some(result), error: None } +} + +fn rpc_error(id: Value, code: i32, message: &str, data: Option) -> RpcResponse { + RpcResponse { + jsonrpc: "2.0", + id, + result: None, + error: Some(RpcError { code, message: message.to_string(), data }), + } +} + +fn rpc_app_error(id: Value, error: AppError) -> RpcResponse { + rpc_error(id, -32603, &error.to_string(), Some(Value::String(error.to_string()))) +} + +fn response_to_string(response: RpcResponse) -> String { + serde_json::to_string(&response).unwrap_or_else(|error| { + json!({ + "jsonrpc": "2.0", + "id": Value::Null, + "error": { + "code": -32603, + "message": "internal error", + "data": error.to_string(), + } + }) + .to_string() + }) +} + +fn search_ast(config: &McpConfig, query: &str) -> Result> { + let search_config = SearchConfig { + queries: vec![query.to_string()], + root_path: config.root_path.clone(), + lang_mode: config.lang_mode.clone(), + }; + let compiled_queries = compile_queries(&search_config)?; + let query_trigram_set = Arc::new(build_query_trigram_set(&search_config.queries)); + let index_path = index_path_for_root(search_config.root_path.as_path()); + let index_manifest = Arc::new(Mutex::new(match load_index(&index_path) { + Ok(manifest) => manifest, + Err(_) => IndexManifest::new(search_config.root_path.clone()), + })); + let outcome = super::run_search( + &search_config, + &compiled_queries, + &query_trigram_set, + &index_manifest, + &ColorMode::Off, + true, + true, + ); + + Ok(outcome.results.into_iter().map(Into::into).collect()) +} + +fn lookup_symbol(config: &McpConfig, name: &str) -> Result> { + if !config.db_path.exists() { + return Err(AppError::DbError(format!( + "no structural index found at {}\n hint: run dora --persist {} first", + config.db_path.display(), + config.root_path.display() + ))); + } + + let db = MemoryDb::open(&config.db_path)?; + let symbols = { + let exact = db.find_symbols_by_name(name)?; + if exact.is_empty() { + db.find_symbols_by_name_contains(name)? + } else { + exact + } + }; + + let mut rows = Vec::new(); + for symbol in symbols { + let file = db.get_file_by_id(symbol.file_id)?.ok_or_else(|| { + AppError::DbError(format!("missing file row for file_id {}", symbol.file_id)) + })?; + rows.push(LookupSymbolItem::from((symbol, file))); + } + + rows.sort_by(|left, right| { + left.file_path + .cmp(&right.file_path) + .then_with(|| left.start_line.cmp(&right.start_line)) + .then_with(|| left.start_col.cmp(&right.start_col)) + .then_with(|| left.name.cmp(&right.name)) + }); + + Ok(rows) +} + +fn compile_queries( + config: &SearchConfig, +) -> Result>>> { + let query = config + .queries + .first() + .ok_or_else(|| AppError::QueryCompileError("query must not be empty".to_string()))?; + + match &config.lang_mode { + LangMode::Single(lang) => { + let ts_lang = super::lang_to_ts_language(lang); + let compiled = query::compile_query(&ts_lang, query)?; + Ok(Arc::new(HashMap::from([( + lang.clone(), + Arc::new(query::MultiCompiledQuery { queries: vec![compiled], language: ts_lang }), + )]))) + } + LangMode::Auto => { + let mut map = HashMap::new(); + for (lang, ts_lang) in get_all_languages() { + if let Ok(compiled) = query::compile_query(&ts_lang, query) { + map.insert( + lang, + Arc::new(query::MultiCompiledQuery { + queries: vec![compiled], + language: ts_lang, + }), + ); + } + } + if map.is_empty() { + return Err(AppError::QueryCompileError(format!( + "query did not compile against any supported language\n query: {}\n hint: check the S-expression syntax and node type names", + query + ))); + } + Ok(Arc::new(map)) + } + } +} + +impl From for SearchAstItem { + fn from(result: MatchResult) -> Self { + Self { + file_path: result.file_path.display().to_string(), + capture_name: result.capture_name, + matched_text: result.matched_text, + start_line: result.start_line, + start_col: result.start_col, + end_line: result.end_line, + end_col: result.end_col, + } + } +} + +impl From<(SymbolRow, FileRow)> for LookupSymbolItem { + fn from(value: (SymbolRow, FileRow)) -> Self { + let (symbol, file) = value; + Self { + file_path: file.path, + kind: symbol.kind.to_string(), + name: symbol.name, + start_line: symbol.start_line, + start_col: symbol.start_col, + end_line: symbol.end_line, + end_col: symbol.end_col, + signature: symbol.signature, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memory::{NewFileRow, NewSymbolRow, SymbolKind}; + use crate::types::Language; + use std::fs; + use tempfile::TempDir; + + fn make_config(root_path: PathBuf, db_path: PathBuf, lang_mode: LangMode) -> McpConfig { + McpConfig { root_path, db_path, lang_mode } + } + + #[test] + fn test_unknown_method_returns_error() { + let config = + make_config(PathBuf::from("."), PathBuf::from("/tmp/missing.db"), LangMode::Auto); + let request = RpcRequest { + jsonrpc: "2.0".to_string(), + id: Some(json!(1)), + method: "nope".to_string(), + params: None, + }; + let response = handle_request(&config, request); + assert_eq!(response.error.unwrap().code, -32601); + } + + #[test] + fn test_missing_params_returns_invalid_params() { + let config = + make_config(PathBuf::from("."), PathBuf::from("/tmp/missing.db"), LangMode::Auto); + let request = RpcRequest { + jsonrpc: "2.0".to_string(), + id: Some(json!(1)), + method: "search_ast".to_string(), + params: None, + }; + let response = handle_request(&config, request); + assert_eq!(response.error.unwrap().code, -32602); + } + + #[test] + fn test_lookup_symbol_exact_then_contains() { + let dir = TempDir::new().unwrap(); + let db_path = dir.path().join("index.db"); + let db = MemoryDb::open(&db_path).unwrap(); + let file_id = db + .upsert_file(&NewFileRow { + path: dir.path().join("a.rs").display().to_string(), + mtime: 1, + language: "rust".to_string(), + }) + .unwrap(); + db.insert_symbol(&NewSymbolRow { + file_id, + kind: SymbolKind::Function, + name: "authenticate".to_string(), + start_line: 1, + start_col: 0, + end_line: 1, + end_col: 12, + signature: Some("fn authenticate()".to_string()), + }) + .unwrap(); + let config = make_config(dir.path().to_path_buf(), db_path, LangMode::Auto); + let exact = lookup_symbol(&config, "authenticate").unwrap(); + assert_eq!(exact.len(), 1); + let fallback = lookup_symbol(&config, "auth").unwrap(); + assert_eq!(fallback.len(), 1); + assert_eq!(fallback[0].name, "authenticate"); + } + + #[test] + fn test_search_ast_returns_match_array() { + let dir = TempDir::new().unwrap(); + let file = dir.path().join("simple.rs"); + fs::write(&file, "fn hello() {}\n").unwrap(); + let config = make_config( + dir.path().to_path_buf(), + dir.path().join("index.db"), + LangMode::Single(Language::Rust), + ); + let results = search_ast(&config, "(function_item name: (identifier) @fn_name)").unwrap(); + assert!(!results.is_empty()); + assert!(results.iter().any(|row| row.matched_text == "hello")); + } +} diff --git a/src/memory.rs b/src/memory.rs index 9991981..e277d56 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -296,6 +296,14 @@ impl MemoryDb { ) } + pub fn find_symbols_by_name_contains(&self, needle: &str) -> Result> { + let pattern = format!("%{needle}%"); + self.query_symbols( + "SELECT id, file_id, kind, name, start_line, start_col, end_line, end_col, signature FROM symbols WHERE name LIKE ?1 ORDER BY name ASC, file_id ASC", + params![pattern], + ) + } + pub fn find_symbols_by_kind(&self, kind: &SymbolKind) -> Result> { self.query_symbols( "SELECT id, file_id, kind, name, start_line, start_col, end_line, end_col, signature FROM symbols WHERE kind = ?1 ORDER BY name ASC, file_id ASC", @@ -548,6 +556,18 @@ mod tests { assert_eq!(symbols.len(), 2); } + #[test] + fn test_find_symbols_by_name_contains() { + let db = make_db(); + let file_id = insert_file(&db, "/tmp/a.rs", 100); + db.insert_symbol(&make_symbol_row(file_id, SymbolKind::Function, "authenticate", 1)) + .unwrap(); + db.insert_symbol(&make_symbol_row(file_id, SymbolKind::Function, "authorise", 2)).unwrap(); + db.insert_symbol(&make_symbol_row(file_id, SymbolKind::Function, "connect", 3)).unwrap(); + let symbols = db.find_symbols_by_name_contains("auth").unwrap(); + assert_eq!(symbols.len(), 2); + } + #[test] fn test_find_symbols_by_kind() { let db = make_db();