mcp-memory/server/src/handlers/workspaces.rs

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::<SnippetsTool>(
            "snippets",
            "Consolidated reusable code snippets management (store, search, delete, tag)",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        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::<Vec<_>>();
                        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::<ReadDirectoryArchitectureTool>(
            "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<MemoryState>) -> crate::error::Result<String> {
        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<String> {
            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::<Result<Vec<_>, 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::<SemanticCodeSearchTool>(
            "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<MemoryState>) -> crate::error::Result<String> {
        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::<ManageSubagentNamespaceTool>(
            "manage_subagent_namespace",
            "Manage isolated memory namespaces for subagent sessions (create, condense/promote, or purge)",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        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<Entity> = 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");
    }
}