mcp-memory/server/src/handlers_v2/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;
use std::time::{SystemTime, UNIX_EPOCH};

pub struct PinFileHandler;

#[async_trait]
impl McpTool for PinFileHandler {
    fn name(&self) -> &'static str {
        "pin_file"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<PinFileTool>("pin_file", "Execute pin_file")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: PinFileTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.pinned_files.modify(|pinned| {
            pinned.retain(|p| p.namespace != req.namespace || p.file_path != req.file_path);
            pinned.push(crate::models::PinnedFile {
                namespace: req.namespace,
                file_path: req.file_path,
                timestamp: SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs(),
                git_branch: req.git_branch,
            });
        });
        Ok("File pinned".to_string())
    }
}

pub struct UnpinFileHandler;

#[async_trait]
impl McpTool for UnpinFileHandler {
    fn name(&self) -> &'static str {
        "unpin_file"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Execute unpin_file")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: UnpinFileTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.pinned_files.modify(|pinned| {
            pinned.retain(|p| p.namespace != req.namespace || p.file_path != req.file_path)
        });
        Ok("File unpinned".to_string())
    }
}

pub struct ListPinnedFilesHandler;

#[async_trait]
impl McpTool for ListPinnedFilesHandler {
    fn name(&self) -> &'static str {
        "list_pinned_files"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<ListPinnedFilesTool>(
            "list_pinned_files",
            "Execute list_pinned_files",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: ListPinnedFilesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state.pinned_files.read_with(|pinned| {
            let filtered: Vec<_> = pinned.iter().filter(|p| {
                let ns_match = match &req.namespace {
                    Some(ns) => &p.namespace == ns,
                    std::option::Option::None => true,
                };
                let branch_match = match &req.git_branch {
                    Some(branch) => p.git_branch.is_none() || p.git_branch.as_deref() == Some(branch.as_str()),
                    std::option::Option::None => true,
                };
                ns_match && branch_match
            }).collect();
            serde_json::to_string(&filtered).unwrap_or_default()
        });
        Ok(data.to_string())
    }
}

pub struct StoreSnippetHandler;

#[async_trait]
impl McpTool for StoreSnippetHandler {
    fn name(&self) -> &'static str {
        "store_snippet"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Execute store_snippet")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: StoreSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let snippet = Snippet {
            name: req.name.clone(),
            language: req.language,
            code: req.code,
            description: req.description,
            updated_at: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs(),
        };

        let s_clone = snippet.clone();
        state.snippets.modify(|snippets| {
            snippets.retain(|s| s.name != req.name);
            snippets.push(s_clone);
        });

        if let Ok(idx) = state.search_index.read() {
            drop(idx.index_snippet(&snippet));
        }

        Ok(format!("Snippet '{}' stored.", req.name).to_string())
    }
}

pub struct SearchSnippetsHandler;

#[async_trait]
impl McpTool for SearchSnippetsHandler {
    fn name(&self) -> &'static str {
        "search_snippets"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<SearchSnippetsTool>("search_snippets", "Execute search_snippets")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: SearchSnippetsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let query = req.query;
        let data = state.snippets.read_with(|snippets| {
            let results: Vec<_> = snippets.iter().filter(|s| {
                contains_ignore_ascii_case(&s.name, &query)
                    || contains_ignore_ascii_case(&s.description, &query)
                    || contains_ignore_ascii_case(&s.language, &query)
            }).cloned().collect();
            serde_json::to_string(&results).unwrap_or_default()
        });
        Ok(data.to_string())
    }
}

pub struct DeleteSnippetHandler;

#[async_trait]
impl McpTool for DeleteSnippetHandler {
    fn name(&self) -> &'static str {
        "delete_snippet"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<DeleteSnippetTool>("delete_snippet", "Execute delete_snippet")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: DeleteSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let mut deleted = false;
        state.snippets.modify(|snippets| {
            let orig = snippets.len();
            snippets.retain(|s| s.name != req.name);
            deleted = snippets.len() < orig;
        });
        if deleted {
            Ok("Snippet deleted.".to_string())
        } else {
            Ok("Snippet not found.".to_string())
        }
    }
}

pub struct SaveContextWorkspaceHandler;

#[async_trait]
impl McpTool for SaveContextWorkspaceHandler {
    fn name(&self) -> &'static str {
        "save_context_workspace"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<SaveContextWorkspaceTool>(
            "save_context_workspace",
            "Execute save_context_workspace",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: SaveContextWorkspaceTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.context_workspaces.modify(|ws| {
            ws.retain(|w| w.namespace != req.namespace || w.name != req.name);
            ws.push(crate::models::ContextWorkspace {
                namespace: req.namespace,
                name: req.name,
                pinned_files: req.pinned_files,
                active_task_ids: req.active_task_ids,
                saved_at: SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs(),
            });
        });
        Ok("Context workspace saved".to_string())
    }
}

pub struct LoadContextWorkspaceHandler;

#[async_trait]
impl McpTool for LoadContextWorkspaceHandler {
    fn name(&self) -> &'static str {
        "load_context_workspace"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<LoadContextWorkspaceTool>(
            "load_context_workspace",
            "Execute load_context_workspace",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: LoadContextWorkspaceTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state.context_workspaces.read_with(|ws| {
            let filtered: Vec<_> = ws.iter().filter(|w| {
                w.namespace == req.namespace && w.name == req.name
            }).collect();
            serde_json::to_string(&filtered.first()).unwrap_or_default()
        });
        Ok(data.to_string())
    }
}

pub struct ListContextWorkspacesHandler;

#[async_trait]
impl McpTool for ListContextWorkspacesHandler {
    fn name(&self) -> &'static str {
        "list_context_workspaces"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<ListContextWorkspacesTool>(
            "list_context_workspaces",
            "Execute list_context_workspaces",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: ListContextWorkspacesTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state.context_workspaces.read_with(|ws| {
            let filtered: Vec<_> = ws.iter().filter(|w| {
                w.namespace == req.namespace
            }).collect();
            serde_json::to_string(&filtered).unwrap_or_default()
        });
        Ok(data.to_string())
    }
}

pub struct AddPrChecklistItemHandler;

#[async_trait]
impl McpTool for AddPrChecklistItemHandler {
    fn name(&self) -> &'static str {
        "add_pr_checklist_item"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<AddPrChecklistItemTool>(
            "add_pr_checklist_item",
            "Execute add_pr_checklist_item",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: AddPrChecklistItemTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.pr_checklists.modify(|items| {
            items.push(crate::models::PrChecklistItem {
                namespace: req.namespace,
                id: uuid::Uuid::new_v4().to_string(),
                description: req.description,
            })
        });
        Ok("PR checklist item added".to_string())
    }
}

pub struct GetPrChecklistHandler;

#[async_trait]
impl McpTool for GetPrChecklistHandler {
    fn name(&self) -> &'static str {
        "get_pr_checklist"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<GetPrChecklistTool>("get_pr_checklist", "Execute get_pr_checklist")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: GetPrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state.pr_checklists.read_with(|items| {
            let filtered: Vec<_> = items.iter().filter(|i| {
                i.namespace == req.namespace
            }).collect();
            serde_json::to_string(&filtered).unwrap_or_default()
        });
        Ok(data.to_string())
    }
}

pub struct ClearPrChecklistHandler;

#[async_trait]
impl McpTool for ClearPrChecklistHandler {
    fn name(&self) -> &'static str {
        "clear_pr_checklist"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<ClearPrChecklistTool>(
            "clear_pr_checklist",
            "Execute clear_pr_checklist",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: ClearPrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        state
            .pr_checklists
            .modify(|items| items.retain(|i| i.namespace != req.namespace));
        Ok("PR checklist cleared".to_string())
    }
}

use crate::handlers_v2::utils::*;