diff --git a/Cargo.lock b/Cargo.lock index e1cfd79..79d2128 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -409,6 +409,16 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "bstr" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bb31b46c14244e20ee9984b11bf5c992b91fb6939fea616e3512c8baecdbe5f" +dependencies = [ + "memchr", + "serde_core", +] + [[package]] name = "built" version = "0.8.1" @@ -1403,6 +1413,19 @@ dependencies = [ "url", ] +[[package]] +name = "globset" +version = "0.4.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07c34a9410465b45bd9787443bc7370f37735bad04b0f0cd57ff1a3186c98988" +dependencies = [ + "aho-corasick", + "bstr", + "log", + "regex-automata", + "regex-syntax", +] + [[package]] name = "h2" version = "0.4.19" @@ -1763,6 +1786,22 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "ignore" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00b69833ed729dc5aa7d19541d96d6cf8e9137194207a04916d658e43168402f" +dependencies = [ + "crossbeam-deque", + "globset", + "log", + "memchr", + "regex-automata", + "same-file", + "walkdir", + "winapi-util", +] + [[package]] name = "image" version = "0.25.10" @@ -2114,6 +2153,7 @@ dependencies = [ "fastembed", "futures-util", "git2", + "ignore", "image", "redb", "reqwest 0.13.5", @@ -3356,6 +3396,15 @@ dependencies = [ "tempfile", ] +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "schannel" version = "0.1.29" @@ -4521,6 +4570,16 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -4686,6 +4745,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" diff --git a/server/Cargo.toml b/server/Cargo.toml index 8551b08..2827363 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -37,6 +37,7 @@ tree-sitter-java = "0.23.5" tree-sitter-c = "0.24.2" tree-sitter-cpp = "0.23.4" tree-sitter-go = "0.25.0" +ignore = "0.4.33" [build-dependencies] chrono = "0.4.45" diff --git a/server/src/indexer.rs b/server/src/indexer.rs new file mode 100644 index 0000000..049445a --- /dev/null +++ b/server/src/indexer.rs @@ -0,0 +1,134 @@ +use crate::state::MemoryState; +use crate::models::Snippet; +use crate::embedding::generate_embeddings_async; +use ignore::WalkBuilder; +use std::sync::Arc; +use tree_sitter::{Node, Parser}; + +pub async fn start_background_indexer(state: Arc) { + // Determine workspace root. Since mcp-memory is usually run from its project root, + // current_dir is a good default for the git repository. + let workspace_root = std::env::current_dir().unwrap_or_else(|_| state.base_dir.clone()); + + tokio::spawn(async move { + tracing::info!("Starting background indexer in {:?}", workspace_root); + + let walker = WalkBuilder::new(&workspace_root) + .hidden(true) + .git_ignore(true) + .build(); + + let mut files_to_process = Vec::new(); + for result in walker { + match result { + Ok(entry) => { + if entry.file_type().map_or(false, |ft| ft.is_file()) { + let path = entry.path().to_path_buf(); + let ext = path.extension().and_then(|e| e.to_str()).unwrap_or(""); + if ["rs", "ts", "js", "jsx", "tsx", "py", "java", "c", "cpp", "go"].contains(&ext) { + files_to_process.push(path); + } + } + } + Err(e) => { + tracing::warn!("Error walking directory: {}", e); + } + } + } + + let idx = state.get_search_index(); + + for file_path in files_to_process { + if let Ok(content) = std::fs::read_to_string(&file_path) { + let ext = file_path.extension().and_then(|e| e.to_str()).unwrap_or(""); + let language = match ext { + "rs" => tree_sitter_rust::LANGUAGE, + "ts" | "tsx" | "js" | "jsx" => tree_sitter_typescript::LANGUAGE_TYPESCRIPT, + "py" => tree_sitter_python::LANGUAGE, + "java" => tree_sitter_java::LANGUAGE, + "c" | "h" => tree_sitter_c::LANGUAGE, + "cpp" | "cc" | "cxx" | "hpp" | "hxx" => tree_sitter_cpp::LANGUAGE, + "go" => tree_sitter_go::LANGUAGE, + _ => continue, + }; + + let mut parser = Parser::new(); + if parser.set_language(&language.into()).is_err() { + continue; + } + + if let Some(tree) = parser.parse(&content, None) { + let mut chunks = Vec::new(); + extract_chunks(tree.root_node(), &content, &mut chunks, ext); + + for (name, code, desc) in chunks { + // Generate embedding + if let Ok(mut emb) = generate_embeddings_async(vec![code.clone()]).await { + let embedding = emb.pop(); + + let file_name = file_path.file_name().unwrap_or_default().to_string_lossy(); + let snippet_name = format!("{}:{}", file_name, name); + + let snippet = Snippet { + name: snippet_name.to_string(), + language: ext.to_string(), + code: code.clone(), + description: format!("{} in {}", desc, file_path.display()), + updated_at: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + embedding, + }; + + state.snippets.modify(|snippets| { + // Prevent duplicates if already indexed + if !snippets.iter().any(|s| s.name == snippet.name) { + snippets.push(snippet.clone()); + } + }); + + let _ = idx.index_snippet(&snippet).await; + } + } + } + } + } + tracing::info!("Background indexing completed."); + }); +} + +fn extract_chunks(node: Node, code: &str, chunks: &mut Vec<(String, String, String)>, ext: &str) { + let kind = node.kind(); + let is_structural = matches!( + kind, + "function_item" + | "function_declaration" + | "function_definition" + | "method_definition" + | "struct_item" + | "class_declaration" + ); + + if is_structural { + let code_text = node.utf8_text(code.as_bytes()).unwrap_or("").to_string(); + + let mut name = "unknown".to_string(); + let mut cursor = node.walk(); + for child in node.children(&mut cursor) { + let child_kind = child.kind(); + if child_kind == "identifier" || child_kind == "name" || child_kind == "type_identifier" { + name = child.utf8_text(code.as_bytes()).unwrap_or("unknown").to_string(); + break; + } + } + + let desc = format!("{} AST node", kind); + chunks.push((name, code_text, desc)); + } else { + let mut cursor = node.walk(); + for child in node.named_children(&mut cursor) { + extract_chunks(child, code, chunks, ext); + } + } +} diff --git a/server/src/main.rs b/server/src/main.rs index 8df44e7..fa8a54c 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -12,6 +12,7 @@ mod models; mod router; mod search; pub mod embedding; +pub mod indexer; mod state; mod store; mod tools; @@ -131,6 +132,10 @@ async fn run_server(state: Arc) -> Result<(), Box