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().is_some_and(|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); } } }