use crate::models::*; use crate::router::McpTool; use crate::state::MemoryState; use crate::tools::*; use async_trait::async_trait; use serde_json::Value; use std::sync::Arc; pub struct SnippetsHandler; #[async_trait] impl McpTool for SnippetsHandler { fn name(&self) -> &'static str { "snippets" } fn schema(&self) -> Value { crate::mcp::tool_def::( "snippets", "Consolidated reusable code snippets management (store, search, delete, tag)", ) } async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { let req: SnippetsTool = serde_json::from_value(args).map_err(|e| e.to_string())?; match req.action { SnippetAction::Store => { let name = req.name.or(req.query).or(req.id).ok_or_else(|| { crate::error::AppError::Internal("Missing required parameter 'name', 'query', or 'id' as snippet name for action 'store'. Next step: Provide snippet name in 'name' or 'query' field and retry.".to_string()) })?; let lang = req.language.unwrap_or_else(|| "text".to_string()); let code = req.code.unwrap_or_default(); let desc = req.description.unwrap_or_default(); let text_to_embed = format!( "Name: {}\nLanguage: {}\nDescription: {}\nCode: {}", name, lang, desc, code ); let embedding = crate::embedding::generate_embedding_async(text_to_embed) .await .ok(); let snippet = Snippet { name: name.clone(), language: lang, code, description: desc, updated_at: crate::handlers::utils::now_secs(), tags: req.tags.unwrap_or_default(), embedding, namespace: req .namespace .unwrap_or_else(crate::models::default_namespace), repo_name: req.repo_name, origin_file: req.origin_file, line_range: req.line_range, }; let idx = state.get_search_index().await; drop(idx.index_snippet(&snippet)); state.code.snippets.modify(|snippets| { snippets.retain(|s| s.name != name); snippets.push(snippet); }); Ok(format!("Snippet '{}' stored.", name)) } SnippetAction::Search => { let query = req.query.or(req.name).unwrap_or_default(); let tags_filter = req.tags.unwrap_or_default(); let is_hybrid = req.hybrid.unwrap_or(false); if is_hybrid { let snippets = state.code.snippets.read_with(|s| s.clone()); let q_tokens: Vec<_> = query .to_lowercase() .split_whitespace() .map(|s| s.to_string()) .collect(); let mut scored: Vec<_> = snippets .into_iter() .filter_map(|s| { let text = format!("{} {} {} {}", s.name, s.description, s.language, s.code) .to_lowercase(); let mut score = 0.0; for tok in &q_tokens { if text.contains(tok) { score += 1.0; } } if !tags_filter.is_empty() { for tag in &tags_filter { if s.tags .iter() .any(|t| t.to_lowercase() == tag.to_lowercase()) { score += 2.0; } } } if score > 0.0 { Some(serde_json::json!({ "name": s.name, "description": s.description, "language": s.language, "tags": s.tags, "score": score, "code": s.code })) } else { None } }) .collect(); scored.sort_by(|a, b| { b["score"] .as_f64() .unwrap_or(0.0) .total_cmp(&a["score"].as_f64().unwrap_or(0.0)) }); scored.truncate(10); Ok(serde_json::to_string_pretty(&scored)?) } else { let data = state.code.snippets.read_with(|snippets| { let results: Vec<_> = snippets .iter() .filter(|s| { let matches_text = query.is_empty() || contains_ignore_ascii_case(&s.name, &query) || contains_ignore_ascii_case(&s.description, &query) || contains_ignore_ascii_case(&s.language, &query) || contains_ignore_ascii_case(&s.code, &query) || s.tags.iter().any(|t| contains_ignore_ascii_case(t, &query)); let matches_tags = tags_filter.is_empty() || tags_filter.iter().any(|tf| { s.tags.iter().any(|st| contains_ignore_ascii_case(st, tf)) }); matches_text && matches_tags }) .take(10) .map(|s| { serde_json::json!({ "name": s.name, "description": s.description, "language": s.language, "tags": s.tags, "code": s.code, }) }) .collect::>(); serde_json::to_string_pretty(&results) })?; Ok(data) } } SnippetAction::Delete => { let id = req.name.or(req.id).or(req.query).ok_or_else(|| { crate::error::AppError::Internal("Missing required parameter 'name', 'id', or 'query' for action 'delete'. Next step: Provide snippet ID/name in request and retry.".to_string()) })?; let mut deleted = false; state.code.snippets.modify(|snippets| { let initial = snippets.len(); snippets.retain(|s| s.name != id); deleted = snippets.len() < initial; }); if deleted { Ok(format!("Snippet '{}' deleted.", id)) } else { Err(crate::error::AppError::Internal(format!( "Snippet '{}' not found.", id ))) } } SnippetAction::Tag => { let id = req.name.or(req.id).or(req.query).ok_or_else(|| { crate::error::AppError::Internal("Missing required parameter 'name', 'id', or 'query' for action 'tag'. Next step: Provide snippet ID/name and 'tags' array in request and retry.".to_string()) })?; let tags = req.tags.unwrap_or_default(); let mut found = false; state.code.snippets.modify(|snippets| { for s in snippets.iter_mut() { if s.name == id { for tag in &tags { if !s.tags.contains(tag) { s.tags.push(tag.clone()); } } found = true; break; } } }); if found { Ok(format!("Snippet '{}' tagged with: {:?}", id, tags)) } else { Err(crate::error::AppError::Internal(format!( "Snippet '{}' not found.", id ))) } } } } } use crate::handlers::utils::*; use crate::tools::ReadDirectoryArchitectureTool; use std::fs; pub struct ReadDirectoryArchitectureHandler; #[async_trait] impl McpTool for ReadDirectoryArchitectureHandler { fn name(&self) -> &'static str { "read_directory_architecture" } fn schema(&self) -> Value { crate::mcp::tool_def::( "read_directory_architecture", "Get a bird's-eye view of a directory, reading the file tree and extracting a basic structural summary.", ) } async fn execute(&self, args: Value, _state: Arc) -> crate::error::Result { let tool_args: ReadDirectoryArchitectureTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let dir = tool_args.directory.clone(); let result = tokio::task::spawn_blocking(move || -> crate::error::Result { let mut summary = String::new(); fn visit_dirs( dir: &std::path::Path, summary: &mut String, depth: usize, ) -> std::io::Result<()> { if depth > 10 { return Ok(()); } if dir.is_dir() { let mut entries = fs::read_dir(dir)?.collect::, std::io::Error>>()?; entries.sort_by_key(|e| e.path()); for entry in entries { let path = entry.path(); let indent = " ".repeat(depth); let name = entry.file_name().to_string_lossy().to_string(); if name.starts_with('.') || name == "target" || name == "node_modules" || name == "dist" { continue; } if path.is_dir() { summary.push_str(&format!("{}- {}/\n", indent, name)); visit_dirs(&path, summary, depth + 1)?; } else { // Extract a brief 1-line heuristic if it's a known file type let mut peek = String::new(); if let Ok(content) = fs::read_to_string(&path) { // Find the first docstring or struct/class definition for line in content.lines() { let t = line.trim(); if t.starts_with("///") || t.starts_with("# ") || t.starts_with("struct ") || t.starts_with("class ") || t.starts_with("function ") { let truncated: String = t.chars().take(80).collect(); peek = format!(" -> {}", truncated); break; } } } summary.push_str(&format!("{}- {}{}\n", indent, name, peek)); } } } Ok(()) } let path = std::path::Path::new(&dir); if !path.exists() { return Err(crate::error::AppError::Internal(format!( "Directory does not exist: {}", dir ))); } summary.push_str(&format!("Architecture of {}:\n", dir)); visit_dirs(path, &mut summary, 0).map_err(|e| e.to_string())?; Ok(summary) }) .await .map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??; Ok(result) } } use crate::tools::SemanticCodeSearchTool; pub struct SemanticCodeSearchHandler; #[async_trait] impl McpTool for SemanticCodeSearchHandler { fn name(&self) -> &'static str { "semantic_code_search" } fn schema(&self) -> Value { crate::mcp::tool_def::( "semantic_code_search", "Perform a semantic vector search across indexed code snippets and knowledge graph nodes using fastembed.", ) } async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { let tool_args: SemanticCodeSearchTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let results = state .search() .semantic_search(&tool_args.query, None, 5) .await?; if results.is_empty() { return Ok(format!( "No semantic matches found for query: {}", tool_args.query )); } let mut out = format!("Semantic Search Results for '{}':\n", tool_args.query); for res in &results { out.push_str(&format!( "- [{:.2}] {}: {}\n", res.score, res.title, res.body )); } if state.ollama.is_available().await { let prompt = format!( "You are an expert software architect. Synthesize a concise GraphRAG summary of the following graph nodes relevant to the query: '{}'.\n\nData:\n{}", tool_args.query, out ); if let Ok(summary) = state .ollama .generate(&prompt, None, Some("Respond clearly and concisely."), None) .await { out.push_str("\n\n--- Local GraphRAG Summary ---\n"); out.push_str(&summary); } } Ok(out) } } pub struct ManageSubagentNamespaceHandler; #[async_trait] impl McpTool for ManageSubagentNamespaceHandler { fn name(&self) -> &'static str { "manage_subagent_namespace" } fn schema(&self) -> Value { crate::mcp::tool_def::( "manage_subagent_namespace", "Manage isolated memory namespaces for subagent sessions (create, condense/promote, or purge)", ) } async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { let req: ManageSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let subagent_ns = format!("subagent-{}", req.subagent_id); match req.action { SubagentNamespaceAction::Create => { state.record_activity( "namespace", &format!("Created subagent namespace {}", subagent_ns), None, ); Ok(format!( "Subagent memory namespace '{}' initialized.", subagent_ns )) } SubagentNamespaceAction::Purge => { state.modify_graph(|g| { g.entities.retain(|_, e| e.namespace != subagent_ns); g.relations.retain(|r| r.namespace != subagent_ns); }); state.record_activity( "namespace", &format!("Purged subagent namespace {}", subagent_ns), None, ); Ok(format!( "Subagent memory namespace '{}' purged.", subagent_ns )) } SubagentNamespaceAction::Condense => { let purge = req.purge_after_promotion.unwrap_or(true); let mut promoted_entities = 0; let mut promoted_relations = 0; state.modify_graph(|g| { let sub_entities: Vec = g .entities .values() .filter(|e| e.namespace == subagent_ns) .cloned() .collect(); for mut e in sub_entities { promoted_entities += 1; e.namespace = "global".to_string(); let global_entry = g.entities.entry(e.name.clone()).or_insert_with(|| Entity { name: e.name.clone(), entity_type: e.entity_type.clone(), observations: Vec::new(), namespace: "global".to_string(), git_branch: e.git_branch.clone(), ..Default::default() }); global_entry.namespace = "global".to_string(); for obs in e.observations { if !global_entry.observations.contains(&obs) { global_entry.observations.push(obs); } } } for r in &mut g.relations { if r.namespace == subagent_ns { promoted_relations += 1; r.namespace = "global".to_string(); } } if purge { g.entities.retain(|_, e| e.namespace != subagent_ns); g.relations.retain(|r| r.namespace != subagent_ns); } }); Ok(format!( "Subagent namespace '{}' condensed: promoted {} entities and {} relations to global graph.", subagent_ns, promoted_entities, promoted_relations )) } } } } #[cfg(test)] mod tests { use super::*; use serde_json::json; use tempfile::tempdir; #[tokio::test] async fn test_snippets_lifecycle() { let dir = tempdir().unwrap(); let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap())); let snippets_handler = SnippetsHandler; // 1. Store let store_args = json!({ "action": "store", "name": "init_db", "language": "sql", "description": "Initialize database", "code": "SELECT 1;", "tags": ["db", "sql"] }); let res1 = snippets_handler .execute(store_args, state.clone()) .await .unwrap(); assert_eq!(res1, "Snippet 'init_db' stored."); // 2. Search let search_args = json!({ "action": "search", "query": "SELECT" }); let res2 = snippets_handler .execute(search_args, state.clone()) .await .unwrap(); assert!(res2.contains("init_db")); // 3. Tag let tag_args = json!({ "action": "tag", "name": "init_db", "tags": ["postgres"] }); let res3 = snippets_handler .execute(tag_args, state.clone()) .await .unwrap(); assert!(res3.contains("tagged with")); // 4. Delete let delete_args = json!({ "action": "delete", "name": "init_db" }); let res4 = snippets_handler .execute(delete_args, state.clone()) .await .unwrap(); assert_eq!(res4, "Snippet 'init_db' deleted."); } #[tokio::test] async fn test_subagent_namespace_lifecycle() { let dir = tempdir().unwrap(); let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap())); let manage_sub_ns = ManageSubagentNamespaceHandler; // Create let create_res = manage_sub_ns .execute( json!({ "subagent_id": "auditor_1", "action": "create" }), state.clone(), ) .await .unwrap(); assert!(create_res.contains("initialized")); // Add subagent entity into graph state.modify_graph(|g| { g.entities.insert( "SubNode".to_string(), Entity { name: "SubNode".to_string(), entity_type: "Module".to_string(), observations: vec!["test observation".to_string()], namespace: "subagent-auditor_1".to_string(), ..Default::default() }, ); }); // Condense let condense_res = manage_sub_ns .execute( json!({ "subagent_id": "auditor_1", "action": "condense", "purge_after_promotion": true }), state.clone(), ) .await .unwrap(); assert!(condense_res.contains("condensed: promoted 1 entities")); // Verify entity is now global let graph = state.graph.cache.read().unwrap(); assert_eq!(graph.entities.get("SubNode").unwrap().namespace, "global"); } }