mcp-memory/server/src/indexer.rs

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<MemoryState>) {
    // 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);
        }
    }
}