use crate::models::*; use crate::state::MemoryState; use crate::tools::*; macro_rules! parse_tool { ($args:expr, $id:expr, $type:ty) => { match parse_args::<$type>($args) { Ok(r) => r, Err(e) => { let response = Some(crate::mcp::success( $id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": format!("Invalid args: {}", e)}] }), )); tracing::trace!("Returning response from handle_request: {:?}", response); return response; } } }; } use serde::de::DeserializeOwned; use std::collections::HashSet; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; fn parse_args(args: serde_json::Value) -> Result { serde_json::from_value(args).map_err(|e| format!("Invalid args: {}", e)) } pub struct MemoryHandler { pub state: Arc, } impl MemoryHandler { pub async fn handle_request(&self, req: serde_json::Value) -> Option { let start_time = std::time::Instant::now(); let id = req.get("id").cloned().unwrap_or(serde_json::Value::Null); let id_clone = id.clone(); let method = req.get("method").and_then(|m| m.as_str()).unwrap_or(""); let tool_name = if method == "tools/call" { req.get("params").and_then(|p| p.get("name")).and_then(|n| n.as_str()).unwrap_or("unknown") } else { "" }; if method == "tools/call" { tracing::info!(">>> [Server] Handling MCP tool call: {} (id: {})", tool_name, id); } else { tracing::debug!(">>> [Server] Handling MCP request method: {}", method); } tracing::trace!(">>> [Server] Full MCP Request payload: {}", req.to_string()); let response = match method { "server/discover" => { let payload = serde_json::json!({ "resultType": "complete", "ttlMs": 0, "cacheScope": "public", "supportedVersions": ["2026-07-28", "2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05"], "capabilities": { "tools": serde_json::json!({}) }, "_meta": { "io.modelcontextprotocol/serverInfo": { "name": "gemini-mcp-memory", "version": "3.0.0" } } }); tracing::debug!( "<<< [Server] Replying to server/discover with payload: {}", payload.to_string() ); Some(crate::mcp::success(id, payload)) } "initialize" => { let init = rmcp::model::InitializeResult::new( rmcp::model::ServerCapabilities::builder() .enable_tools() .build(), ) .with_server_info(rmcp::model::Implementation::new( "gemini-mcp-memory", "3.0.0", )); tracing::debug!("<<< [Server] Replying to initialize with rmcp payload"); Some(crate::mcp::success( id, serde_json::to_value(&init).unwrap(), )) } "notifications/initialized" => None, "tools/list" => { let tools = vec![ crate::mcp::tool_def::( "query_graph_path", "Traverse the knowledge graph to find a path between two entities.", ), crate::mcp::tool_def::( "create_entities", "Create new entities in the knowledge graph.", ), crate::mcp::tool_def::( "create_relations", "Create new relations between entities in the knowledge graph.", ), crate::mcp::tool_def::( "add_observations", "Add new observations to existing entities in the knowledge graph.", ), crate::mcp::tool_def::( "delete_entities", "Delete entities from the knowledge graph.", ), crate::mcp::tool_def::( "delete_observations", "Delete observations from existing entities.", ), crate::mcp::tool_def::( "delete_relations", "Delete relations between entities.", ), crate::mcp::tool_def::( "read_graph", "Read the entire knowledge graph.", ), crate::mcp::tool_def::( "search_nodes", "Search for entities in the knowledge graph by name or type.", ), crate::mcp::tool_def::( "open_nodes", "Open and retrieve full details of specific nodes in the knowledge graph.", ), crate::mcp::tool_def::( "log_code_change", "Log a significant code change or refactor in the memory system.", ), crate::mcp::tool_def::( "query_recent_changes", "Query recently logged code changes.", ), crate::mcp::tool_def::( "visualize_graph", "Generate a visual representation of the knowledge graph.", ), crate::mcp::tool_def::( "add_sticky_note", "Add a sticky note for unstructured thoughts or reminders.", ), crate::mcp::tool_def::( "read_sticky_notes", "Read all active sticky notes.", ), crate::mcp::tool_def::( "delete_sticky_note", "Delete a specific sticky note by its 1-indexed position.", ), crate::mcp::tool_def::( "clear_sticky_notes", "Clear all active sticky notes.", ), crate::mcp::tool_def::( "condense_entity", "Condense or summarize an entity's observations to reduce size.", ), crate::mcp::tool_def::( "add_task", "Add a new task to the task tracker.", ), crate::mcp::tool_def::( "update_task_status", "Update the status of an existing task.", ), crate::mcp::tool_def::( "delete_task", "Delete a task and all its children.", ), crate::mcp::tool_def::( "list_active_tasks", "List all currently active tasks.", ), crate::mcp::tool_def::( "set_acceptance_criteria", "Define a strict checklist of acceptance criteria for a given task.", ), crate::mcp::tool_def::( "verify_acceptance_criteria", "Mark a previously defined acceptance criteria as met.", ), crate::mcp::tool_def::( "store_snippet", "Store a reusable code snippet.", ), crate::mcp::tool_def::( "search_snippets", "Search through stored code snippets.", ), crate::mcp::tool_def::( "delete_snippet", "Delete a stored code snippet.", ), crate::mcp::tool_def::( "log_decision", "Log an architectural decision record (ADR).", ), crate::mcp::tool_def::( "query_decisions", "Query architectural decision records.", ), crate::mcp::tool_def::( "merge_entities", "Merge two entities in the knowledge graph into one.", ), crate::mcp::tool_def::( "find_orphans", "Find orphaned entities (entities without any relations) in the graph.", ), crate::mcp::tool_def::( "learn_preference", "Record a user preference or behavior to adapt future interactions.", ), crate::mcp::tool_def::( "read_preferences", "Read all learned user preferences.", ), crate::mcp::tool_def::( "log_error_fix", "Log a complex error and its fix for future reference.", ), crate::mcp::tool_def::( "search_error_fixes", "Search through previously logged error fixes.", ), crate::mcp::tool_def::( "pin_file", "Pin a file to keep it explicitly in the context workspace.", ), crate::mcp::tool_def::( "unpin_file", "Unpin a file from the context workspace.", ), crate::mcp::tool_def::( "list_pinned_files", "List all currently pinned files.", ), crate::mcp::tool_def::( "add_session_summary", "Add a summary of the current session.", ), crate::mcp::tool_def::( "get_project_timeline", "Get a timeline of major project events.", ), crate::mcp::tool_def::( "leave_handoff_memo", "Leave a memo for the next session or agent.", ), crate::mcp::tool_def::( "read_handoff_memos", "Read pending handoff memos.", ), crate::mcp::tool_def::( "clear_handoff_memos", "Clear handoff memos after reading.", ), crate::mcp::tool_def::( "update_env_fingerprint", "Update the environment fingerprint (e.g., OS, tool versions).", ), crate::mcp::tool_def::( "read_env_fingerprint", "Read the current environment fingerprint.", ), crate::mcp::tool_def::( "log_env_requirement", "Log a required tool or package for the environment.", ), crate::mcp::tool_def::( "add_milestone", "Add a new project milestone.", ), crate::mcp::tool_def::( "update_milestone", "Update the status of a project milestone.", ), crate::mcp::tool_def::( "list_milestones", "List all project milestones.", ), crate::mcp::tool_def::( "generate_standup_report", "Generate a standup report summarizing recent work, blockers, and next steps.", ), crate::mcp::tool_def::( "register_environment", "Register details about a specific deployment environment.", ), crate::mcp::tool_def::( "get_environment_details", "Get detailed information about a specific deployment environment.", ), crate::mcp::tool_def::( "add_pr_checklist_item", "Add an item to the PR checklist.", ), crate::mcp::tool_def::( "get_pr_checklist", "Get the current PR checklist.", ), crate::mcp::tool_def::( "clear_pr_checklist", "Clear the PR checklist.", ), crate::mcp::tool_def::( "log_tech_debt", "Log identified technical debt.", ), crate::mcp::tool_def::( "resolve_tech_debt", "Mark a logged technical debt as resolved.", ), crate::mcp::tool_def::( "list_tech_debt", "List all unresolved technical debt.", ), crate::mcp::tool_def::( "save_context_workspace", "Save the current set of pinned files and context.", ), crate::mcp::tool_def::( "load_context_workspace", "Load a previously saved context workspace.", ), crate::mcp::tool_def::( "list_context_workspaces", "List all saved context workspaces.", ), crate::mcp::tool_def::( "omni_search", "Search across all memory sources (graph, tasks, snippets, ADRs, etc.) at once.", ), crate::mcp::tool_def::( "get_project_health", "Get a synthesized health report of the project based on memory data.", ), ]; Some(crate::mcp::success( id, serde_json::json!({ "tools": tools }), )) } "tools/call" => { let params = req.get("params").unwrap_or(&serde_json::Value::Null); let name = params.get("name").and_then(|n| n.as_str()).unwrap_or(""); let args = params .get("arguments") .cloned() .unwrap_or(serde_json::Value::Object(Default::default())); self.state .broadcast_activity(&format!("Agent executed tool: {}", name)); let result: Result = match name { "query_graph_path" => { let req = parse_tool!(args.clone(), id, crate::tools::QueryGraphPathTool); let graph = self.state.get_full_graph(); let max_depth = req.max_depth.unwrap_or(5); let mut queue = std::collections::VecDeque::new(); let mut visited = std::collections::HashSet::new(); let mut parents: std::collections::HashMap = std::collections::HashMap::new(); queue.push_back(req.start_node.clone()); visited.insert(req.start_node.clone()); let mut found = false; let mut current_depth = 0; let mut nodes_at_current_depth = 1; let mut nodes_at_next_depth = 0; while let Some(current) = queue.pop_front() { if current == req.end_node { found = true; break; } nodes_at_current_depth -= 1; if current_depth < max_depth { for rel in &graph.relations { if rel.from == current && !visited.contains(&rel.to) { visited.insert(rel.to.clone()); parents.insert( rel.to.clone(), (current.clone(), rel.relation_type.clone()), ); queue.push_back(rel.to.clone()); nodes_at_next_depth += 1; } else if rel.to == current && !visited.contains(&rel.from) { visited.insert(rel.from.clone()); parents.insert( rel.from.clone(), ( current.clone(), format!("inverse({})", rel.relation_type), ), ); queue.push_back(rel.from.clone()); nodes_at_next_depth += 1; } } } if nodes_at_current_depth == 0 { current_depth += 1; nodes_at_current_depth = nodes_at_next_depth; nodes_at_next_depth = 0; } } if found { let mut path = Vec::new(); let mut curr = req.end_node.clone(); while curr != req.start_node { let (parent, rel) = parents.get(&curr).unwrap(); path.push(format!("({}) --[{}]--> ({})", parent, rel, curr)); curr = parent.clone(); } path.reverse(); Ok(format!("Path found:\n{}", path.join("\n"))) } else { Ok(format!( "No path found between {} and {} within depth {}", req.start_node, req.end_node, max_depth )) } } "create_entities" => { let req = parse_tool!(args.clone(), id, CreateEntitiesTool); self.state.modify_graph(|g| { for entity in req.entities { if !entity.name.is_empty() { if let Ok(idx) = self.state.search_index.read() { drop(idx.index_entity(&entity)); } g.entities.insert(entity.name.clone(), entity); } } }); Ok("Entities created".to_string()) } "create_relations" => { let req = parse_tool!(args.clone(), id, CreateRelationsTool); self.state.modify_graph(|g| { for relation in req.relations { if !relation.from.is_empty() && !relation.to.is_empty() { g.relations.push(relation); } } }); Ok("Relations created".to_string()) } "add_observations" => { let req = parse_tool!(args.clone(), id, AddObservationsTool); let full = self.state.get_full_graph(); self.state.modify_graph(|g| { for o in req.observations { if let Some(full_e) = full.entities.get(&o.entity_name) { let mut e = g.entities.get(&o.entity_name).cloned().unwrap_or_else( || Entity { name: o.entity_name.clone(), entity_type: full_e.entity_type.clone(), observations: vec![], namespace: full_e.namespace.clone(), git_branch: None, }, ); e.observations.extend(o.contents); g.entities.insert(o.entity_name, e); } } }); Ok("Observations added".to_string()) } "delete_entities" => { let req = parse_tool!(args.clone(), id, DeleteEntitiesTool); let to_delete: HashSet<_> = req.entity_names.into_iter().collect(); self.state.modify_graph(|master| { for name in &to_delete { master.entities.remove(name); } master.relations.retain(|r| { !to_delete.contains(&r.from) && !to_delete.contains(&r.to) }); }); Ok("Entities deleted".to_string()) } "delete_observations" => { let req = parse_tool!(args.clone(), id, DeleteObservationsTool); self.state.modify_graph(|master| { for d in req.deletions { if let Some(e) = master.entities.get_mut(&d.entity_name) { let to_rem: HashSet<_> = d.observations.into_iter().collect(); e.observations.retain(|o| !to_rem.contains(o)); } } }); Ok("Observations deleted".to_string()) } "delete_relations" => { let req = parse_tool!(args.clone(), id, DeleteRelationsTool); self.state.modify_graph(|master| { let to_rem: HashSet<_> = req.relations.into_iter().collect(); master.relations.retain(|r| !to_rem.contains(r)); }); Ok("Relations deleted".to_string()) } "read_graph" => { let req = parse_tool!(args.clone(), id, ReadGraphTool); let mut full = self.state.get_full_graph(); if let Some(ns) = req.namespace { full.entities.retain(|_, e| e.namespace == ns); full.relations.retain(|r| r.namespace == ns); } let data = serde_json::to_string(&full).unwrap_or_default(); Ok(data.to_string()) } "search_nodes" => { let req = parse_tool!(args.clone(), id, SearchNodesTool); let matches = if let Ok(idx) = self.state.search_index.read() { idx.search(&req.query, req.namespace.as_deref()) .unwrap_or_default() } else { vec![] }; let mut result = KnowledgeGraph::default(); let full = self.state.get_full_graph(); for (id, doc_type, _, _, _) in matches { if doc_type == "entity" && let Some(e) = full.entities.get(&id) { result.entities.insert(id, e.clone()); } } let data = serde_json::to_string(&result).unwrap_or_default(); Ok(data.to_string()) } "open_nodes" => { let req = parse_tool!(args.clone(), id, OpenNodesTool); let targets: HashSet<_> = req.names.into_iter().collect(); let full = self.state.get_full_graph(); let mut result = KnowledgeGraph::default(); let mut connected = HashSet::new(); for r in &full.relations { if targets.contains(&r.from) { connected.insert(r.to.clone()); result.relations.push(r.clone()); } else if targets.contains(&r.to) { connected.insert(r.from.clone()); result.relations.push(r.clone()); } } for (name, e) in full.entities { if targets.contains(&name) || connected.contains(&name) { result.entities.insert(name, e); } } let data = serde_json::to_string(&result).unwrap_or_default(); Ok(data.to_string()) } "log_code_change" => { let req = parse_tool!(args.clone(), id, LogCodeChangeTool); self.state.ledger.modify(|ledger| { ledger.push(CodeChange { timestamp: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), file_path: req.file_path, description: req.description, git_commit: req.git_commit, git_branch: req.git_branch, }); }); Ok("Code change logged".to_string()) } "query_recent_changes" => { let data = serde_json::to_string(&self.state.ledger.read()) .unwrap_or_else(|_| "[]".to_string()); Ok(data.to_string()) } "visualize_graph" => { let req = parse_tool!(args.clone(), id, VisualizeGraphTool); let query = req.query.unwrap_or_default().to_lowercase(); let full = self.state.get_full_graph(); let mut included = HashSet::new(); for (name, e) in &full.entities { if let Some(ns) = &req.namespace && e.namespace != *ns { continue; } if query.is_empty() || name.to_lowercase().contains(&query) || e.entity_type.to_lowercase().contains(&query) { included.insert(name.clone()); } } let mut to_draw = Vec::new(); for r in &full.relations { if let Some(ns) = &req.namespace && r.namespace != *ns { continue; } if query.is_empty() || included.contains(&r.from) || included.contains(&r.to) { included.insert(r.from.clone()); included.insert(r.to.clone()); to_draw.push(r); } } let mut output = String::from("graph TD;\n"); for name in &included { let safe = name.replace("\"", "").replace("(", "").replace(")", ""); output.push_str(&format!( " id_{}[\"{}\"];\n", safe.replace(" ", "_").replace("-", "_").replace(".", "_"), safe )); } for r in to_draw { let safe_f = r .from .replace(" ", "_") .replace("-", "_") .replace(".", "_") .replace("\"", "") .replace("(", "") .replace(")", ""); let safe_t = r.to.replace(" ", "_") .replace("-", "_") .replace(".", "_") .replace("\"", "") .replace("(", "") .replace(")", ""); let safe_r = r.relation_type.replace("\"", ""); output.push_str(&format!( " id_{}-->|\"{}\"|id_{};\n", safe_f, safe_r, safe_t )); } if output == "graph TD;\n" { output = "No nodes found to visualize.".to_string(); } Ok(output.to_string()) } "add_sticky_note" => { let req = parse_tool!(args.clone(), id, AddStickyNoteTool); self.state.sticky.modify(|notes| { notes.push(StickyNote { timestamp: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), content: req.content, }); }); Ok("Sticky note added.".to_string()) } "read_sticky_notes" => { let data = serde_json::to_string(&self.state.sticky.read()) .unwrap_or_else(|_| "[]".to_string()); Ok(data.to_string()) } "delete_sticky_note" => { let req = parse_tool!(args.clone(), id, DeleteStickyNoteTool); let mut success = false; self.state.sticky.modify(|notes| { if req.index > 0 && req.index <= notes.len() { notes.remove(req.index - 1); success = true; } }); if success { Ok("Sticky note deleted.".to_string()) } else { Err("Invalid sticky note index.".to_string()) } } "clear_sticky_notes" => { self.state.sticky.modify(|notes| { notes.clear(); }); Ok("All sticky notes cleared.".to_string()) } "condense_entity" => { let req = parse_tool!(args.clone(), id, CondenseEntityTool); self.state.modify_graph(|master| { if let Some(e) = master.entities.get_mut(&req.entity_name) { e.observations = req.summarized_observations; } }); Ok("Entity condensed".to_string()) } "add_task" => { let req = parse_tool!(args.clone(), id, AddTaskTool); let now = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(); let task_id = uuid::Uuid::new_v4().to_string(); let parent_id = req.parent_id.clone(); let deps = req.dependencies.clone().unwrap_or_default(); let task = Task { id: task_id.clone(), title: req.title, status: "pending".to_string(), description: req.description, created_at: now, updated_at: now, git_branch: req.git_branch, parent_id, dependencies: deps, acceptance_criteria: vec![], }; if let Ok(idx) = self.state.search_index.read() { drop(idx.index_task(&task)); } self.state.tasks.modify(|tasks| { tasks.push(task); }); Ok(format!("Task added with ID: {}", task_id).to_string()) } "delete_task" => { let req = parse_tool!(args.clone(), id, DeleteTaskTool); let mut deleted_count = 0; self.state.tasks.modify(|tasks| { let initial_len = tasks.len(); // Collect IDs of tasks to delete (this task + all its recursive children) let mut to_delete = std::collections::HashSet::new(); to_delete.insert(req.id.clone()); let mut added_new = true; while added_new { added_new = false; for t in tasks.iter() { if let Some(pid) = &t.parent_id && to_delete.contains(pid) && !to_delete.contains(&t.id) { to_delete.insert(t.id.clone()); added_new = true; } } } tasks.retain(|t| !to_delete.contains(&t.id)); deleted_count = initial_len - tasks.len(); }); if deleted_count > 0 { Ok(vec![ format!("Deleted task and its children ({} total).", deleted_count) .to_string(), ][0] .clone()) } else { Ok("Task not found.".to_string()) } } "update_task_status" => { let req = parse_tool!(args.clone(), id, UpdateTaskStatusTool); let mut found = false; let mut blocked = false; let mut blocker_details = String::new(); let target_status = req.status.to_lowercase(); self.state.tasks.modify(|tasks| { // Find target task let mut target_id = String::new(); if let Some(t) = tasks.iter().find(|t| t.id == req.id || t.title == req.id) { target_id = t.id.clone(); } if target_id.is_empty() { return; } found = true; if target_status == "done" || target_status == "completed" { // 1. Check Acceptance Criteria if let Some(t) = tasks.iter().find(|t| t.id == target_id) && t.acceptance_criteria.iter().any(|c| !c.is_met) { blocked = true; blocker_details = "Unmet acceptance criteria exist.".to_string(); } // 2. Check dependencies if !blocked { let mut uncompleted_deps = Vec::new(); if let Some(t) = tasks.iter().find(|t| t.id == target_id) { for dep_id in &t.dependencies { if let Some(dep_task) = tasks.iter().find(|dt| dt.id == *dep_id) && dep_task.status != "completed" && dep_task.status != "done" { uncompleted_deps.push(dep_task.title.clone()); } } } if !uncompleted_deps.is_empty() { blocked = true; blocker_details = format!( "Blocked by dependencies: {}", uncompleted_deps.join(", ") ); } } // 3. Check child tasks if !blocked { let mut uncompleted_children = Vec::new(); for child in tasks .iter() .filter(|t| t.parent_id.as_ref() == Some(&target_id)) { if child.status != "completed" && child.status != "done" { uncompleted_children.push(child.title.clone()); } } if !uncompleted_children.is_empty() { blocked = true; blocker_details = format!( "Blocked by child tasks: {}", uncompleted_children.join(", ") ); } } } if !blocked { // Apply update if let Some(t) = tasks.iter_mut().find(|t| t.id == target_id) { t.status = target_status.clone(); t.updated_at = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(); } // Cascade cancellation to children if target_status == "cancelled" || target_status == "abandoned" { let mut to_cancel = vec![target_id.clone()]; let mut i = 0; while i < to_cancel.len() { let current_pid = to_cancel[i].clone(); for t in tasks.iter_mut() { if t.parent_id.as_ref() == Some(¤t_pid) && t.status != "completed" { t.status = target_status.clone(); to_cancel.push(t.id.clone()); } } i += 1; } } } }); if blocked { Ok(vec![format!( "Error: Cannot transition task. {}", blocker_details )] .into_iter() .next() .unwrap()) } else if found { Ok("Task status updated.".to_string()) } else { Ok("Task not found.".to_string()) } } "list_active_tasks" => { let req = parse_tool!(args.clone(), id, ListActiveTasksTool); let mut tasks = self.state.tasks.read(); tasks.retain(|t| t.status != "done"); if let Some(branch) = req.git_branch { tasks.retain(|t| { t.git_branch.is_none() || t.git_branch.as_deref() == Some(branch.as_str()) }); } let data = serde_json::to_string(&tasks).unwrap_or_default(); Ok(data.to_string()) } "set_acceptance_criteria" => { let req = parse_tool!(args.clone(), id, SetAcceptanceCriteriaTool); let mut success = false; self.state.tasks.modify(|tasks| { if let Some(task) = tasks.iter_mut().rev().find(|t| t.title == req.task_title) { task.acceptance_criteria = req .criteria .into_iter() .map(|desc| crate::models::AcceptanceCriteria { id: uuid::Uuid::new_v4().to_string(), description: desc, is_met: false, }) .collect(); task.updated_at = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(); success = true; } }); if success { Ok( vec!["Acceptance criteria set successfully.".to_string()][0] .clone(), ) } else { Ok("Task not found.".to_string()) } } "verify_acceptance_criteria" => { let req = parse_tool!(args.clone(), id, VerifyAcceptanceCriteriaTool); let mut success = false; let mut already_met = false; self.state.tasks.modify(|tasks| { if let Some(task) = tasks.iter_mut().find(|t| t.id == req.task_id) && let Some(ac) = task .acceptance_criteria .iter_mut() .find(|c| c.id == req.criteria || c.description == req.criteria) { if ac.is_met { already_met = true; } else { ac.is_met = true; success = true; task.updated_at = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(); } } }); if success { Ok(vec![format!( "Acceptance criteria verified with proof: {}", req.proof )][0] .clone()) } else if already_met { Ok("Acceptance criteria was already met.".to_string()) } else { Ok( vec!["Acceptance criteria or task not found.".to_string()][0] .clone(), ) } } "store_snippet" => { let req = parse_tool!(args.clone(), id, StoreSnippetTool); 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() .as_secs(), }; let s_clone = snippet.clone(); self.state.snippets.modify(|snippets| { snippets.retain(|s| s.name != req.name); snippets.push(s_clone); }); if let Ok(idx) = self.state.search_index.read() { drop(idx.index_snippet(&snippet)); } Ok(format!("Snippet '{}' stored.", req.name).to_string()) } "search_snippets" => { let req = parse_tool!(args.clone(), id, SearchSnippetsTool); let query = req.query.to_lowercase(); let snippets = self.state.snippets.read(); let mut results = Vec::new(); for s in snippets { if s.name.to_lowercase().contains(&query) || s.description.to_lowercase().contains(&query) || s.language.to_lowercase().contains(&query) { results.push(s); } } let data = serde_json::to_string(&results).unwrap_or_default(); Ok(data.to_string()) } "delete_snippet" => { let req = parse_tool!(args.clone(), id, DeleteSnippetTool); let mut deleted = false; self.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()) } } "log_decision" => { let req = parse_tool!(args.clone(), id, LogDecisionTool); let mut adr_id = String::new(); let mut new_adr = None; self.state.adrs.modify(|adrs| { adr_id = format!("ADR-{:04}", adrs.len() + 1); let a = Adr { id: adr_id.clone(), title: req.title, context: req.context, decision: req.decision, consequence: req.consequence, timestamp: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), }; new_adr = Some(a.clone()); adrs.push(a); }); if let Some(adr) = new_adr { if let Ok(idx) = self.state.search_index.read() { drop(idx.index_adr(&adr)); } } Ok(format!("Decision logged as {}", adr_id).to_string()) } "query_decisions" => { let req = parse_tool!(args.clone(), id, QueryDecisionsTool); let mut adrs = self.state.adrs.read(); if let Some(q) = req.query { let q = q.to_lowercase(); adrs.retain(|a| { a.title.to_lowercase().contains(&q) || a.context.to_lowercase().contains(&q) || a.decision.to_lowercase().contains(&q) }); } let data = serde_json::to_string(&adrs).unwrap_or_default(); Ok(data.to_string()) } "merge_entities" => { let req = parse_tool!(args.clone(), id, MergeEntitiesTool); self.state.modify_graph(|master| { if let Some(src) = master.entities.remove(&req.source_entity) { if let Some(tgt) = master.entities.get_mut(&req.target_entity) { tgt.observations.extend(src.observations); tgt.observations = MemoryState::unique_items(tgt.observations.clone()); } else { let mut new_tgt = src.clone(); new_tgt.name = req.target_entity.clone(); master.entities.insert(req.target_entity.clone(), new_tgt); } } for r in &mut master.relations { if r.from == req.source_entity { r.from = req.target_entity.clone(); } if r.to == req.source_entity { r.to = req.target_entity.clone(); } } master.relations = MemoryState::unique_items(master.relations.clone()); }); Ok("Entities merged".to_string()) } "find_orphans" => { let full = self.state.get_full_graph(); let mut connected = std::collections::HashSet::new(); for r in &full.relations { connected.insert(r.from.clone()); connected.insert(r.to.clone()); } let orphans: Vec = full .entities .keys() .filter(|k| !connected.contains(*k)) .cloned() .collect(); let data = serde_json::to_string(&orphans).unwrap_or_default(); Ok(data.to_string()) } "learn_preference" => { let req = parse_tool!(args.clone(), id, LearnPreferenceTool); self.state.prefs.modify(|prefs| { prefs.insert( req.key.clone(), crate::models::Preference { key: req.key.clone(), value: req.value, updated_at: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), }, ); }); Ok("Preference learned".to_string()) } "read_preferences" => { let prefs = self.state.prefs.read(); let data = serde_json::to_string(&prefs).unwrap_or_default(); Ok(data.to_string()) } "log_error_fix" => { let req = parse_tool!(args.clone(), id, LogErrorFixTool); self.state.error_fixes.modify(|fixes| { fixes.push(crate::models::ErrorFix { signature: req.signature, solution: req.solution, timestamp: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), git_commit: req.git_commit, git_branch: req.git_branch, }) }); Ok("Error fix logged".to_string()) } "search_error_fixes" => { let req = parse_tool!(args.clone(), id, SearchErrorFixesTool); let q = req.query.to_lowercase(); let mut fixes = self.state.error_fixes.read(); fixes.retain(|f| { f.signature.to_lowercase().contains(&q) || f.solution.to_lowercase().contains(&q) }); let data = serde_json::to_string(&fixes).unwrap_or_default(); Ok(data.to_string()) } "pin_file" => { let req = parse_tool!(args.clone(), id, PinFileTool); self.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() .as_secs(), git_branch: req.git_branch, }); }); Ok("File pinned".to_string()) } "unpin_file" => { let req = parse_tool!(args.clone(), id, UnpinFileTool); self.state.pinned_files.modify(|pinned| { pinned.retain(|p| { !(p.namespace == req.namespace && p.file_path == req.file_path) }) }); Ok("File unpinned".to_string()) } "list_pinned_files" => { let req = parse_tool!(args.clone(), id, ListPinnedFilesTool); let mut pinned = self.state.pinned_files.read(); if let Some(ns) = req.namespace { pinned.retain(|p| p.namespace == ns); } if let Some(branch) = req.git_branch { pinned.retain(|p| { p.git_branch.is_none() || p.git_branch.as_deref() == Some(branch.as_str()) }); } let data = serde_json::to_string(&pinned).unwrap_or_default(); Ok(data.to_string()) } "add_session_summary" => { let req = parse_tool!(args.clone(), id, AddSessionSummaryTool); self.state.session_summaries.modify(|summaries| { summaries.push(crate::models::SessionSummary { summary: req.summary, namespace: req.namespace, timestamp: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), }) }); Ok("Session summary added".to_string()) } "get_project_timeline" => { let req = parse_tool!(args.clone(), id, GetProjectTimelineTool); let mut summaries = self.state.session_summaries.read(); if let Some(ns) = req.namespace { summaries.retain(|s| s.namespace == ns); } summaries.sort_by_key(|s| s.timestamp); let data = serde_json::to_string(&summaries).unwrap_or_default(); Ok(data.to_string()) } "leave_handoff_memo" => { let req = parse_tool!(args.clone(), id, LeaveHandoffMemoTool); self.state.handoff_memos.modify(|memos| { memos.push(crate::models::HandoffMemo { id: uuid::Uuid::new_v4().to_string(), author: "agy".to_string(), content: req.content, namespace: req.namespace, timestamp: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), }) }); Ok("Handoff memo left".to_string()) } "read_handoff_memos" => { let req = parse_tool!(args.clone(), id, ReadHandoffMemosTool); let mut memos = self.state.handoff_memos.read(); if let Some(ns) = req.namespace { memos.retain(|m| m.namespace == ns); } let data = serde_json::to_string(&memos).unwrap_or_default(); Ok(data.to_string()) } "clear_handoff_memos" => { let req = parse_tool!(args.clone(), id, ClearHandoffMemosTool); let ids: HashSet<_> = req.ids.into_iter().collect(); self.state .handoff_memos .modify(|memos| memos.retain(|m| !ids.contains(&m.id))); Ok("Handoff memos cleared".to_string()) } "update_env_fingerprint" => { let req = parse_tool!(args.clone(), id, UpdateEnvFingerprintTool); self.state.env_fingerprints.modify(|fps| { fps.insert( req.namespace.clone(), crate::models::EnvFingerprint { namespace: req.namespace.clone(), os: std::env::consts::OS.to_string(), shell: std::env::var("SHELL") .unwrap_or_else(|_| "unknown".to_string()), tool_versions: req.tool_versions, updated_at: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), }, ); }); Ok("Env fingerprint updated".to_string()) } "read_env_fingerprint" => { let req = parse_tool!(args.clone(), id, ReadEnvFingerprintTool); let fps = self.state.env_fingerprints.read(); if let Some(fp) = fps.get(&req.namespace) { let data = serde_json::to_string(fp).unwrap_or_default(); Ok(data.to_string()) } else { Ok("{}".to_string()) } } "log_env_requirement" => { let req = parse_tool!(args.clone(), id, LogEnvRequirementTool); self.state.env_requirements.modify(|reqs| { reqs.retain(|r| !(r.namespace == req.namespace && r.key == req.key)); reqs.push(crate::models::EnvRequirement { namespace: req.namespace, key: req.key, description: req.description, is_secret: req.is_secret, }); }); Ok("Env requirement logged".to_string()) } "add_milestone" => { let req = parse_tool!(args.clone(), id, AddMilestoneTool); self.state.milestones.modify(|ms| { ms.push(crate::models::Milestone { id: uuid::Uuid::new_v4().to_string(), title: req.title, status: "pending".to_string(), namespace: req.namespace, target_date: None, }) }); Ok("Milestone added".to_string()) } "update_milestone" => { let req = parse_tool!(args.clone(), id, UpdateMilestoneTool); let mut found = false; self.state.milestones.modify(|ms| { for m in ms.iter_mut() { if m.id == req.id { m.status = req.status.clone(); found = true; break; } } }); if found { Ok("Milestone updated".to_string()) } else { Ok("Milestone not found".to_string()) } } "list_milestones" => { let req = parse_tool!(args.clone(), id, ListMilestonesTool); let mut ms = self.state.milestones.read(); if let Some(ns) = req.namespace { ms.retain(|m| m.namespace == ns); } let data = serde_json::to_string(&ms).unwrap_or_default(); Ok(data.to_string()) } "generate_standup_report" => { let req = parse_tool!(args.clone(), id, GenerateStandupReportTool); let cutoff = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs() .saturating_sub(req.hours_lookback * 3600); let tasks = self .state .tasks .read() .into_iter() .filter(|t| t.updated_at >= cutoff) .collect::>(); let changes = self .state .ledger .read() .into_iter() .filter(|c| c.timestamp >= cutoff) .collect::>(); let summaries = self .state .session_summaries .read() .into_iter() .filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff) .collect::>(); let report = serde_json::json!({ "tasks_updated": tasks, "code_changes": changes, "session_summaries": summaries }); Ok(report.to_string()) } "register_environment" => { let req = parse_tool!(args.clone(), id, RegisterEnvironmentTool); self.state.environments.modify(|envs| { envs.retain(|e| !(e.namespace == req.namespace && e.name == req.name)); envs.push(crate::models::EnvironmentDetail { namespace: req.namespace, name: req.name, url: req.url, description: req.description, requires_vpn: req.requires_vpn, updated_at: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), }); }); Ok("Environment registered".to_string()) } "get_environment_details" => { let req = parse_tool!(args.clone(), id, GetEnvironmentDetailsTool); let mut envs = self.state.environments.read(); envs.retain(|e| e.namespace == req.namespace); let data = serde_json::to_string(&envs).unwrap_or_default(); Ok(data.to_string()) } "add_pr_checklist_item" => { let req = parse_tool!(args.clone(), id, AddPrChecklistItemTool); self.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()) } "get_pr_checklist" => { let req = parse_tool!(args.clone(), id, GetPrChecklistTool); let mut items = self.state.pr_checklists.read(); items.retain(|i| i.namespace == req.namespace); let data = serde_json::to_string(&items).unwrap_or_default(); Ok(data.to_string()) } "clear_pr_checklist" => { let req = parse_tool!(args.clone(), id, ClearPrChecklistTool); self.state .pr_checklists .modify(|items| items.retain(|i| i.namespace != req.namespace)); Ok("PR checklist cleared".to_string()) } "log_tech_debt" => { let req = parse_tool!(args.clone(), id, LogTechDebtTool); self.state.tech_debts.modify(|debts| { debts.push(crate::models::TechDebt { id: uuid::Uuid::new_v4().to_string(), namespace: req.namespace, description: req.description, ideal_solution: req.ideal_solution, is_resolved: false, created_at: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), git_commit: req.git_commit, git_branch: req.git_branch, }) }); Ok("Tech debt logged".to_string()) } "resolve_tech_debt" => { let req = parse_tool!(args.clone(), id, ResolveTechDebtTool); let mut found = false; self.state.tech_debts.modify(|debts| { for d in debts.iter_mut() { if d.id == req.id { d.is_resolved = true; found = true; break; } } }); if found { Ok("Tech debt resolved".to_string()) } else { Ok("Tech debt not found".to_string()) } } "list_tech_debt" => { let req = parse_tool!(args.clone(), id, ListTechDebtTool); let mut debts = self.state.tech_debts.read(); debts.retain(|d| { d.namespace == req.namespace && (req.include_resolved || !d.is_resolved) }); let data = serde_json::to_string(&debts).unwrap_or_default(); Ok(data.to_string()) } "save_context_workspace" => { let req = parse_tool!(args.clone(), id, SaveContextWorkspaceTool); self.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() .as_secs(), }); }); Ok("Context workspace saved".to_string()) } "load_context_workspace" => { let req = parse_tool!(args.clone(), id, LoadContextWorkspaceTool); let mut ws = self.state.context_workspaces.read(); ws.retain(|w| w.namespace == req.namespace && w.name == req.name); let data = serde_json::to_string(&ws.first()).unwrap_or_default(); Ok(data.to_string()) } "list_context_workspaces" => { let req = parse_tool!(args.clone(), id, ListContextWorkspacesTool); let mut ws = self.state.context_workspaces.read(); ws.retain(|w| w.namespace == req.namespace); let data = serde_json::to_string(&ws).unwrap_or_default(); Ok(data.to_string()) } "omni_search" => { let req = parse_tool!(args.clone(), id, OmniSearchTool); let matches = if let Ok(idx) = self.state.search_index.read() { idx.search(&req.query, req.namespace.as_deref()) .unwrap_or_default() } else { vec![] }; let mut kg = KnowledgeGraph::default(); let mut tasks = Vec::new(); let mut snippets = Vec::new(); let mut adrs = Vec::new(); let full = self.state.get_full_graph(); for (id, doc_type, _, _, _) in &matches { if doc_type == "entity" && let Some(e) = full.entities.get(id) { kg.entities.insert(id.clone(), e.clone()); } } for t in self.state.tasks.read() { if matches.iter().any(|(id, typ, _, _, _)| id == &t.id && typ == "task") { tasks.push(t); } } for s in self.state.snippets.read() { if matches .iter() .any(|(id, typ, _, _, _)| id == &s.name && typ == "snippet") { snippets.push(s); } } for a in self.state.adrs.read() { if matches.iter().any(|(id, typ, _, _, _)| id == &a.id && typ == "adr") { adrs.push(a); } } let q = req.query.to_lowercase(); let tech_debts: Vec<_> = self .state .tech_debts .read() .into_iter() .filter(|d| { (req.namespace.is_none() || d.namespace == *req.namespace.as_ref().unwrap()) && (d.description.to_lowercase().contains(&q) || d.ideal_solution.to_lowercase().contains(&q)) }) .collect(); let memos: Vec<_> = self .state .handoff_memos .read() .into_iter() .filter(|m| { (req.namespace.is_none() || m.namespace == *req.namespace.as_ref().unwrap()) && m.content.to_lowercase().contains(&q) }) .collect(); let error_fixes: Vec<_> = self .state .error_fixes .read() .into_iter() .filter(|f| { f.signature.to_lowercase().contains(&q) || f.solution.to_lowercase().contains(&q) }) .collect(); let report = serde_json::json!({ "knowledge_graph": kg.entities, "tasks": tasks, "snippets": snippets, "adrs": adrs, "tech_debts": tech_debts, "handoff_memos": memos, "error_fixes": error_fixes }); Ok(report.to_string()) } "get_project_health" => { let req = parse_tool!(args.clone(), id, GetProjectHealthTool); let active_tasks = self .state .tasks .read() .into_iter() .filter(|t| t.status != "done") .count(); let unresolved_debt = self .state .tech_debts .read() .into_iter() .filter(|d| d.namespace == req.namespace && !d.is_resolved) .count(); let unread_memos = self .state .handoff_memos .read() .into_iter() .filter(|m| m.namespace == req.namespace) .count(); let active_milestones = self .state .milestones .read() .into_iter() .filter(|m| m.namespace == req.namespace && m.status != "done") .count(); let remaining_checklists = self .state .pr_checklists .read() .into_iter() .filter(|c| c.namespace == req.namespace) .count(); let report = serde_json::json!({ "active_tasks": active_tasks, "unresolved_tech_debt": unresolved_debt, "unread_handoff_memos": unread_memos, "active_milestones": active_milestones, "remaining_pr_checklist_items": remaining_checklists }); Ok(report.to_string()) } _ => Err(format!("Unknown tool: {}", name)), }; match result { Ok(text) => Some(crate::mcp::success( id, serde_json::json!({ "content": [{ "type": "text", "text": text }] }), )), Err(e) => Some(crate::mcp::success( id, serde_json::json!({ "isError": true, "content": [{ "type": "text", "text": e }] }), )), } } _ => { if id != serde_json::Value::Null { Some(crate::mcp::error(id, -32601, "Method not found")) } else { None } } }; let elapsed = start_time.elapsed(); if method == "tools/call" { let is_error = response.as_ref().is_some_and(|r| r.get("error").is_some() || r.get("result").and_then(|res| res.get("isError")).and_then(|e| e.as_bool()).unwrap_or(false)); tracing::info!("<<< [Server] MCP tool call {} (id: {}) completed in {:?} [Error: {}]", tool_name, id_clone, elapsed, is_error); // Broadcast completion latency to the UI Activity Feed let status_msg = if is_error { "with error" } else { "successfully" }; self.state.broadcast_activity(&format!("Tool {} completed {} in {:?}", tool_name, status_msg, elapsed)); } else { tracing::debug!("<<< [Server] MCP request method {} completed in {:?}", method, elapsed); } tracing::trace!("Returning response from handle_request: {:?}", response); response } } #[cfg(test)] mod tests { use super::*; use crate::state::MemoryState; use serde_json::json; use std::sync::Arc; #[tokio::test] async fn test_handle_initialize() { let store_dir = std::env::temp_dir().join(format!( "mcp_test_handlers_{}", std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_secs() )); std::fs::create_dir_all(&store_dir).unwrap(); let redb_path = store_dir.join("mcp_store.redb"); let db = Arc::new(redb::Database::create(&redb_path).unwrap()); { let write_txn = db.begin_write().unwrap(); let _ = write_txn.open_table(crate::store::STORE_TABLE); write_txn.commit().unwrap(); } let state = Arc::new(MemoryState { base_dir: store_dir.clone(), graph: crate::store::Store::new("knowledge_graph_master", db.clone()), search_index: std::sync::RwLock::new( crate::search::MemoryIndex::new(&store_dir).unwrap(), ), ledger: crate::store::Store::new("audit_ledger", db.clone()), sticky: crate::store::Store::new("sticky_notes", db.clone()), tasks: crate::store::Store::new("tasks", db.clone()), snippets: crate::store::Store::new("snippets", db.clone()), adrs: crate::store::Store::new("adrs", db.clone()), prefs: crate::store::Store::new("preferences", db.clone()), error_fixes: crate::store::Store::new("error_fixes", db.clone()), pinned_files: crate::store::Store::new("pinned_files", db.clone()), session_summaries: crate::store::Store::new("session_summaries", db.clone()), handoff_memos: crate::store::Store::new("handoff_memos", db.clone()), env_fingerprints: crate::store::Store::new("env_fingerprints", db.clone()), env_requirements: crate::store::Store::new("env_requirements", db.clone()), milestones: crate::store::Store::new("milestones", db.clone()), environments: crate::store::Store::new("environments", db.clone()), pr_checklists: crate::store::Store::new("pr_checklists", db.clone()), tech_debts: crate::store::Store::new("tech_debts", db.clone()), gates: crate::store::Store::new("gates", db.clone()), context_workspaces: crate::store::Store::new("context_workspaces", db.clone()), activity_tx: tokio::sync::broadcast::channel(100).0, }); let handler = MemoryHandler { state }; let req = json!({ "jsonrpc": "2.0", "id": 1, "method": "initialize", "params": { "protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": { "name": "test-client", "version": "1.0.0" } } }); let response = handler .handle_request(req) .await .expect("Expected a response"); assert_eq!(response["id"], 1); assert!(response.get("result").is_some()); let result = &response["result"]; // assert_eq!(result["protocolVersion"], "2024-11-05"); // CRITICAL BUG FIX CHECK: capabilities MUST contain an empty tools object // Note: Currently it is set to `{}` which may cause proxy dropping tools. Let's verify it matches the actual behavior. assert_eq!(result["capabilities"], serde_json::json!({"tools": {}})); assert_eq!(result["serverInfo"]["name"], "gemini-mcp-memory"); } fn setup_test_handler(test_name: &str) -> MemoryHandler { let store_dir = std::env::temp_dir().join(format!( "mcp_test_handlers_{}_{}", test_name, uuid::Uuid::new_v4() )); std::fs::create_dir_all(&store_dir).unwrap(); let redb_path = store_dir.join("mcp_store.redb"); let db = Arc::new(redb::Database::create(&redb_path).unwrap()); { let write_txn = db.begin_write().unwrap(); let _ = write_txn.open_table(crate::store::STORE_TABLE); write_txn.commit().unwrap(); } let state = Arc::new(MemoryState { base_dir: store_dir.clone(), graph: crate::store::Store::new("knowledge_graph_master", db.clone()), search_index: std::sync::RwLock::new( crate::search::MemoryIndex::new(&store_dir).unwrap(), ), ledger: crate::store::Store::new("audit_ledger", db.clone()), sticky: crate::store::Store::new("sticky_notes", db.clone()), tasks: crate::store::Store::new("tasks", db.clone()), snippets: crate::store::Store::new("snippets", db.clone()), adrs: crate::store::Store::new("adrs", db.clone()), prefs: crate::store::Store::new("preferences", db.clone()), error_fixes: crate::store::Store::new("error_fixes", db.clone()), pinned_files: crate::store::Store::new("pinned_files", db.clone()), session_summaries: crate::store::Store::new("session_summaries", db.clone()), handoff_memos: crate::store::Store::new("handoff_memos", db.clone()), env_fingerprints: crate::store::Store::new("env_fingerprints", db.clone()), env_requirements: crate::store::Store::new("env_requirements", db.clone()), milestones: crate::store::Store::new("milestones", db.clone()), environments: crate::store::Store::new("environments", db.clone()), pr_checklists: crate::store::Store::new("pr_checklists", db.clone()), tech_debts: crate::store::Store::new("tech_debts", db.clone()), gates: crate::store::Store::new("gates", db.clone()), context_workspaces: crate::store::Store::new("context_workspaces", db.clone()), activity_tx: tokio::sync::broadcast::channel(100).0, }); MemoryHandler { state } } #[tokio::test] async fn test_handle_tools_list() { let handler = setup_test_handler("tools_list"); let req = json!({ "jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {} }); let response = handler .handle_request(req) .await .expect("Expected a response"); assert_eq!(response["id"], 2); let tools = response["result"]["tools"] .as_array() .expect("Tools must be an array"); assert!(!tools.is_empty()); // Verify a specific tool is registered let add_task_tool = tools .iter() .find(|t| t["name"] == "add_task") .expect("add_task tool missing"); assert_eq!( add_task_tool["description"], "Add a new task to the task tracker." ); } #[tokio::test] async fn test_handle_add_task() { let handler = setup_test_handler("add_task"); let req = json!({ "jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": { "name": "add_task", "arguments": { "title": "Fix bug in handlers", "description": "The proxy drops capabilities.", "git_branch": "master" } } }); let response = handler .handle_request(req) .await .expect("Expected a response"); assert_eq!(response["id"], 3); let content = &response["result"]["content"][0]; assert_eq!(content["type"], "text"); assert!( content["text"] .as_str() .unwrap() .starts_with("Task added with ID: ") ); // Verify task was actually added to store let tasks = handler.state.tasks.read(); assert_eq!(tasks.len(), 1); assert_eq!(tasks[0].title, "Fix bug in handlers"); assert_eq!(tasks[0].status, "pending"); } #[tokio::test] async fn test_handle_create_entities() { let handler = setup_test_handler("create_entities"); let req = json!({ "jsonrpc": "2.0", "id": 4, "method": "tools/call", "params": { "name": "create_entities", "arguments": { "entities": [ { "name": "MemoryHandler", "entityType": "struct", "observations": ["Handles MCP requests natively"], "namespace": "core" } ] } } }); let response = handler .handle_request(req) .await .expect("Expected a response"); assert_eq!(response["id"], 4); let content = &response["result"]["content"][0]; assert_eq!(content["text"], "Entities created"); // Verify entity was actually added to state let session_graph = handler.state.graph.read(); let entity = session_graph .entities .get("MemoryHandler") .expect("Entity should be in session graph"); assert_eq!(entity.entity_type, "struct"); assert_eq!(entity.observations, vec!["Handles MCP requests natively"]); assert_eq!(entity.namespace, "core".to_string()); } #[tokio::test] async fn test_handle_store_snippet() { let handler = setup_test_handler("store_snippet"); let req = json!({ "jsonrpc": "2.0", "id": 5, "method": "tools/call", "params": { "name": "store_snippet", "arguments": { "name": "Test Snippet", "description": "A snippet used for testing", "language": "rust", "code": "fn main() { println!(\"Hello, World!\"); }" } } }); let response = handler .handle_request(req) .await .expect("Expected a response"); assert_eq!(response["id"], 5); let snippets = handler.state.snippets.read(); assert_eq!(snippets.len(), 1); assert_eq!(snippets[0].name, "Test Snippet"); assert_eq!(snippets[0].language, "rust"); } #[tokio::test] async fn test_handle_add_sticky_note() { let handler = setup_test_handler("add_sticky_note"); let req = json!({ "jsonrpc": "2.0", "id": 6, "method": "tools/call", "params": { "name": "add_sticky_note", "arguments": { "content": "Don't forget to check coverage!" } } }); let response = handler .handle_request(req) .await .expect("Expected a response"); assert_eq!(response["id"], 6); let notes = handler.state.sticky.read(); assert_eq!(notes.len(), 1); assert_eq!(notes[0].content, "Don't forget to check coverage!"); } #[tokio::test] async fn test_handle_create_relations() { let handler = setup_test_handler("create_relations"); let req = json!({ "jsonrpc": "2.0", "id": 7, "method": "tools/call", "params": { "name": "create_relations", "arguments": { "relations": [ { "from": "NodeA", "to": "NodeB", "relationType": "depends_on", "namespace": "core" } ] } } }); let response = handler.handle_request(req).await.unwrap(); assert_eq!(response["id"], 7); let session = handler.state.graph.read(); assert_eq!(session.relations.len(), 1); assert_eq!(session.relations[0].from, "NodeA"); assert_eq!(session.relations[0].to, "NodeB"); } #[tokio::test] async fn test_handle_add_observations() { let handler = setup_test_handler("add_observations"); // Pre-populate entity handler.state.graph.modify(|session| { session.entities.insert( "NodeA".to_string(), crate::models::Entity { name: "NodeA".to_string(), entity_type: "class".to_string(), observations: vec!["Initial".to_string()], namespace: "".to_string(), git_branch: None, }, ); }); let req = json!({ "jsonrpc": "2.0", "id": 8, "method": "tools/call", "params": { "name": "add_observations", "arguments": { "observations": [ { "entityName": "NodeA", "contents": ["New observation"] } ] } } }); let _ = handler.handle_request(req).await.unwrap(); let session = handler.state.graph.read(); let entity = session.entities.get("NodeA").unwrap(); assert_eq!(entity.observations, vec!["Initial", "New observation"]); } #[tokio::test] async fn test_handle_delete_entities() { let handler = setup_test_handler("delete_entities"); handler.state.graph.modify(|session| { session.entities.insert( "ToDelete".to_string(), crate::models::Entity { name: "ToDelete".to_string(), entity_type: "var".to_string(), observations: vec![], namespace: "".to_string(), git_branch: None, }, ); }); // Force flush session to master handler.state.modify_graph(|_| {}); let req = json!({ "jsonrpc": "2.0", "id": 9, "method": "tools/call", "params": { "name": "delete_entities", "arguments": { "entityNames": ["ToDelete"] } } }); let _ = handler.handle_request(req).await.unwrap(); let full_graph = handler.state.get_full_graph(); assert!(full_graph.entities.get("ToDelete").is_none()); } #[tokio::test] async fn test_handle_delete_observations() { let handler = setup_test_handler("delete_observations"); handler.state.graph.modify(|session| { session.entities.insert( "NodeA".to_string(), crate::models::Entity { name: "NodeA".to_string(), entity_type: "class".to_string(), observations: vec!["Keep".to_string(), "Drop".to_string()], namespace: "".to_string(), git_branch: None, }, ); }); handler.state.modify_graph(|_| {}); let req = json!({ "jsonrpc": "2.0", "id": 10, "method": "tools/call", "params": { "name": "delete_observations", "arguments": { "deletions": [ { "entityName": "NodeA", "observations": ["Drop"] } ] } } }); let _ = handler.handle_request(req).await.unwrap(); let full = handler.state.get_full_graph(); let entity = full.entities.get("NodeA").unwrap(); assert_eq!(entity.observations, vec!["Keep"]); } #[tokio::test] async fn test_handle_log_code_change() { let handler = setup_test_handler("log_code_change"); let req = json!({ "jsonrpc": "2.0", "id": 11, "method": "tools/call", "params": { "name": "log_code_change", "arguments": { "filePath": "server/src/handlers.rs", "description": "Added some unit tests", "git_commit": "1234567" } } }); let response = handler.handle_request(req).await.unwrap(); assert_eq!(response["id"], 11); let ledger = handler.state.ledger.read(); assert_eq!(ledger.len(), 1); assert_eq!(ledger[0].file_path, "server/src/handlers.rs"); assert_eq!(ledger[0].git_commit.as_deref(), Some("1234567")); } #[tokio::test] async fn test_handle_list_active_tasks() { let handler = setup_test_handler("list_active_tasks"); handler.state.tasks.modify(|tasks| { tasks.push(crate::models::Task { id: "1".to_string(), title: "Active Task".to_string(), status: "pending".to_string(), description: "".to_string(), created_at: 0, updated_at: 0, git_branch: None, acceptance_criteria: vec![], dependencies: vec![], parent_id: None, }); tasks.push(crate::models::Task { id: "2".to_string(), title: "Completed Task".to_string(), status: "done".to_string(), description: "".to_string(), created_at: 0, updated_at: 0, git_branch: None, acceptance_criteria: vec![], dependencies: vec![], parent_id: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 12, "method": "tools/call", "params": { "name": "list_active_tasks", "arguments": {} } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("Active Task")); assert!(!content.contains("Completed Task")); } #[tokio::test] async fn test_handle_search_snippets() { let handler = setup_test_handler("search_snippets"); handler.state.snippets.modify(|snippets| { snippets.push(crate::models::Snippet { name: "React hook".to_string(), language: "typescript".to_string(), code: "useMemo(() => {}, [])".to_string(), description: "React memoization".to_string(), updated_at: 0, }); snippets.push(crate::models::Snippet { name: "Rust struct".to_string(), language: "rust".to_string(), code: "struct A {}".to_string(), description: "Rust code".to_string(), updated_at: 0, }); }); let req = json!({ "jsonrpc": "2.0", "id": 13, "method": "tools/call", "params": { "name": "search_snippets", "arguments": { "query": "React" } } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("React hook")); assert!(!content.contains("Rust struct")); } #[tokio::test] async fn test_handle_read_sticky_notes() { let handler = setup_test_handler("read_sticky_notes"); handler.state.sticky.modify(|sticky| { sticky.push(crate::models::StickyNote { content: "Remember to commit".to_string(), timestamp: 0, }); }); let req = json!({ "jsonrpc": "2.0", "id": 14, "method": "tools/call", "params": { "name": "read_sticky_notes", "arguments": {} } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("Remember to commit")); } #[tokio::test] async fn test_handle_delete_relations() { let handler = setup_test_handler("delete_relations"); handler.state.graph.modify(|session| { session.relations.push(crate::models::Relation { from: "A".to_string(), to: "B".to_string(), relation_type: "calls".to_string(), namespace: "".to_string(), }); }); handler.state.modify_graph(|_| {}); let req = json!({ "jsonrpc": "2.0", "id": 15, "method": "tools/call", "params": { "name": "delete_relations", "arguments": { "relations": [ { "from": "A", "to": "B", "relationType": "calls", "namespace": "" } ] } } }); let _ = handler.handle_request(req).await.unwrap(); let full = handler.state.get_full_graph(); assert!(full.relations.is_empty()); } #[tokio::test] async fn test_handle_read_graph() { let handler = setup_test_handler("read_graph"); handler.state.graph.modify(|session| { session.entities.insert( "NodeA".to_string(), crate::models::Entity { name: "NodeA".to_string(), entity_type: "var".to_string(), observations: vec![], namespace: "".to_string(), git_branch: None, }, ); }); handler.state.modify_graph(|_| {}); let req = json!({ "jsonrpc": "2.0", "id": 16, "method": "tools/call", "params": { "name": "read_graph", "arguments": {} } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("NodeA")); } #[tokio::test] async fn test_handle_open_nodes() { let handler = setup_test_handler("open_nodes"); let entity = crate::models::Entity { name: "UserRepository".to_string(), entity_type: "class".to_string(), observations: vec!["Handles user data".to_string()], namespace: "".to_string(), git_branch: None, }; handler.state.graph.modify(|session| { session .entities .insert("UserRepository".to_string(), entity); }); handler.state.modify_graph(|_| {}); let req = json!({ "jsonrpc": "2.0", "id": 17, "method": "tools/call", "params": { "name": "open_nodes", "arguments": { "names": ["UserRepository"] } } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("UserRepository")); assert!(content.contains("Handles user data")); } #[tokio::test] async fn test_handle_log_decision() { let handler = setup_test_handler("log_decision"); let req = json!({ "jsonrpc": "2.0", "id": 20, "method": "tools/call", "params": { "name": "log_decision", "arguments": { "title": "Use async I/O", "context": "Need better throughput", "decision": "Use tokio", "consequence": "Requires async all the way down" } } }); let _ = handler.handle_request(req).await.unwrap(); let adrs = handler.state.adrs.read(); assert_eq!(adrs.len(), 1); assert_eq!(adrs[0].title, "Use async I/O"); } #[tokio::test] async fn test_handle_query_decisions() { let handler = setup_test_handler("query_decisions"); handler.state.adrs.modify(|adrs| { adrs.push(crate::models::Adr { id: "adr-1".to_string(), title: "Use PostgreSQL".to_string(), context: "Need relational data".to_string(), decision: "Use pg".to_string(), consequence: "Maintenance overhead".to_string(), timestamp: 0, }); }); let req = json!({ "jsonrpc": "2.0", "id": 21, "method": "tools/call", "params": { "name": "query_decisions", "arguments": { "query": "Postgre" } } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("PostgreSQL")); } #[tokio::test] async fn test_handle_log_error_fix() { let handler = setup_test_handler("log_error_fix"); let req = json!({ "jsonrpc": "2.0", "id": 22, "method": "tools/call", "params": { "name": "log_error_fix", "arguments": { "signature": "IndexOutOfBounds", "solution": "Check array length" } } }); let _ = handler.handle_request(req).await.unwrap(); let fixes = handler.state.error_fixes.read(); assert_eq!(fixes.len(), 1); assert_eq!(fixes[0].signature, "IndexOutOfBounds"); } #[tokio::test] async fn test_handle_search_error_fixes() { let handler = setup_test_handler("search_error_fixes"); handler.state.error_fixes.modify(|fixes| { fixes.push(crate::models::ErrorFix { signature: "NullPointerException".to_string(), solution: "Initialize the pointer".to_string(), timestamp: 0, git_branch: None, git_commit: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 23, "method": "tools/call", "params": { "name": "search_error_fixes", "arguments": { "query": "NullPointer" } } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("Initialize the pointer")); } #[tokio::test] async fn test_handle_list_pinned_files() { let handler = setup_test_handler("list_pinned_files"); handler.state.pinned_files.modify(|files| { files.push(crate::models::PinnedFile { file_path: "src/important.rs".to_string(), timestamp: 0, namespace: "".to_string(), git_branch: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 24, "method": "tools/call", "params": { "name": "list_pinned_files", "arguments": {} } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("important.rs")); } #[tokio::test] async fn test_handle_add_session_summary() { let handler = setup_test_handler("add_session_summary"); let req = json!({ "jsonrpc": "2.0", "id": 25, "method": "tools/call", "params": { "name": "add_session_summary", "arguments": { "namespace": "", "summary": "Finished writing tests" } } }); let _ = handler.handle_request(req).await.unwrap(); let summaries = handler.state.session_summaries.read(); assert_eq!(summaries.len(), 1); assert_eq!(summaries[0].summary, "Finished writing tests"); } #[tokio::test] async fn test_handle_get_project_timeline() { let handler = setup_test_handler("get_project_timeline"); handler.state.session_summaries.modify(|summaries| { summaries.push(crate::models::SessionSummary { summary: "Day 1: Setup project".to_string(), namespace: "".to_string(), timestamp: 0, }); }); let req = json!({ "jsonrpc": "2.0", "id": 26, "method": "tools/call", "params": { "name": "get_project_timeline", "arguments": { "namespace": "" } } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("Day 1: Setup project")); } #[tokio::test] async fn test_handle_log_tech_debt() { let handler = setup_test_handler("log_tech_debt"); let req = json!({ "jsonrpc": "2.0", "id": 27, "method": "tools/call", "params": { "name": "log_tech_debt", "arguments": { "namespace": "", "description": "Hardcoded values", "ideal_solution": "Remove magic numbers" } } }); let _ = handler.handle_request(req).await.unwrap(); let debt = handler.state.tech_debts.read(); assert_eq!(debt.len(), 1); assert_eq!(debt[0].description, "Hardcoded values"); } #[tokio::test] async fn test_handle_list_tech_debt() { let handler = setup_test_handler("list_tech_debt"); handler.state.tech_debts.modify(|debts| { debts.push(crate::models::TechDebt { id: "debt-1".to_string(), description: "Bad naming".to_string(), ideal_solution: "Rename x to num_elements".to_string(), namespace: "".to_string(), is_resolved: false, created_at: 0, git_branch: None, git_commit: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 28, "method": "tools/call", "params": { "name": "list_tech_debt", "arguments": { "namespace": "", "include_resolved": false } } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("Bad naming")); } #[tokio::test] async fn test_handle_get_project_health() { let handler = setup_test_handler("get_project_health"); let req = json!({ "jsonrpc": "2.0", "id": 29, "method": "tools/call", "params": { "name": "get_project_health", "arguments": { "namespace": "" } } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("\"active_tasks\"")); assert!(content.contains("\"unresolved_tech_debt\"")); } #[tokio::test] async fn test_handle_resolve_tech_debt() { let handler = setup_test_handler("resolve_tech_debt"); handler.state.tech_debts.modify(|debts| { debts.push(crate::models::TechDebt { id: "debt-2".to_string(), description: "Old api".to_string(), ideal_solution: "Use new api".to_string(), namespace: "".to_string(), is_resolved: false, created_at: 0, git_branch: None, git_commit: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 30, "method": "tools/call", "params": { "name": "resolve_tech_debt", "arguments": { "id": "debt-2" } } }); let _ = handler.handle_request(req).await.unwrap(); let debts = handler.state.tech_debts.read(); assert!(debts[0].is_resolved); } #[tokio::test] async fn test_handle_leave_handoff_memo() { let handler = setup_test_handler("leave_handoff_memo"); let req = json!({ "jsonrpc": "2.0", "id": 31, "method": "tools/call", "params": { "name": "leave_handoff_memo", "arguments": { "namespace": "", "content": "Make sure to check the logs.", "author": "Riz" } } }); let _ = handler.handle_request(req).await.unwrap(); let memos = handler.state.handoff_memos.read(); assert_eq!(memos.len(), 1); assert_eq!(memos[0].content, "Make sure to check the logs."); } #[tokio::test] async fn test_handle_query_recent_changes() { let handler = setup_test_handler("query_recent_changes"); handler.state.ledger.modify(|ledger| { ledger.push(crate::models::CodeChange { timestamp: 0, file_path: "src/main.rs".to_string(), description: "Fix bug".to_string(), git_commit: None, git_branch: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 18, "method": "tools/call", "params": { "name": "query_recent_changes", "arguments": {} } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("Fix bug")); assert!(content.contains("src/main.rs")); } #[tokio::test] async fn test_handle_update_task_status() { let handler = setup_test_handler("update_task_status"); handler.state.tasks.modify(|tasks| { tasks.push(crate::models::Task { id: "test-task-123".to_string(), title: "In progress task".to_string(), status: "pending".to_string(), description: "".to_string(), created_at: 0, updated_at: 0, git_branch: None, acceptance_criteria: vec![], dependencies: vec![], parent_id: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 19, "method": "tools/call", "params": { "name": "update_task_status", "arguments": { "id": "test-task-123", "status": "in_progress" } } }); let _ = handler.handle_request(req).await.unwrap(); let tasks = handler.state.tasks.read(); assert_eq!(tasks[0].status, "in_progress"); } #[tokio::test] async fn test_handle_pin_file() { let handler = setup_test_handler("pin_file"); let req = json!({ "jsonrpc": "2.0", "id": 100, "method": "tools/call", "params": { "name": "pin_file", "arguments": { "file_path": "/path/to/pinned.rs", "namespace": "global" } } }); let _ = handler.handle_request(req).await.unwrap(); let pinned = handler.state.pinned_files.read(); assert!(pinned.iter().any(|p| p.file_path == "/path/to/pinned.rs")); } #[tokio::test] async fn test_handle_unpin_file() { let handler = setup_test_handler("unpin_file"); handler.state.pinned_files.modify(|files| { files.push(crate::models::PinnedFile { namespace: "global".to_string(), file_path: "/path/to/unpin.rs".to_string(), timestamp: 0, git_branch: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 101, "method": "tools/call", "params": { "name": "unpin_file", "arguments": { "file_path": "/path/to/unpin.rs", "namespace": "global" } } }); let _ = handler.handle_request(req).await.unwrap(); let pinned = handler.state.pinned_files.read(); assert!(!pinned.iter().any(|p| p.file_path == "/path/to/unpin.rs")); } #[tokio::test] async fn test_handle_learn_preference() { let handler = setup_test_handler("learn_preference"); let req = json!({ "jsonrpc": "2.0", "id": 102, "method": "tools/call", "params": { "name": "learn_preference", "arguments": { "key": "formatting", "value": "use spaces" } } }); let _ = handler.handle_request(req).await.unwrap(); let prefs = handler.state.prefs.read(); assert!(prefs.values().any(|p| p.key == "formatting" && p.value == "use spaces")); } #[tokio::test] async fn test_handle_read_preferences() { let handler = setup_test_handler("read_preferences"); handler.state.prefs.modify(|prefs| { prefs.insert("theme".to_string(), crate::models::Preference { key: "theme".to_string(), value: "dark".to_string(), updated_at: 0, }); }); let req = json!({ "jsonrpc": "2.0", "id": 103, "method": "tools/call", "params": { "name": "read_preferences", "arguments": {} } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("theme")); assert!(content.contains("dark")); } #[tokio::test] async fn test_handle_delete_task() { let handler = setup_test_handler("delete_task"); handler.state.tasks.modify(|tasks| { tasks.push(crate::models::Task { id: "task-to-delete".to_string(), title: "".to_string(), status: "".to_string(), description: "".to_string(), created_at: 0, updated_at: 0, git_branch: None, acceptance_criteria: vec![], dependencies: vec![], parent_id: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 104, "method": "tools/call", "params": { "name": "delete_task", "arguments": { "id": "task-to-delete" } } }); let _ = handler.handle_request(req).await.unwrap(); let tasks = handler.state.tasks.read(); assert!(tasks.is_empty()); } #[tokio::test] async fn test_handle_add_milestone() { let handler = setup_test_handler("add_milestone"); let req = json!({ "jsonrpc": "2.0", "id": 105, "method": "tools/call", "params": { "name": "add_milestone", "arguments": { "title": "v1.0", "namespace": "global" } } }); let _ = handler.handle_request(req).await.unwrap(); let milestones = handler.state.milestones.read(); assert_eq!(milestones.len(), 1); assert_eq!(milestones[0].title, "v1.0"); } #[tokio::test] async fn test_handle_update_milestone() { let handler = setup_test_handler("update_milestone"); handler.state.milestones.modify(|ms| { ms.push(crate::models::Milestone { id: "v1.0".to_string(), title: "v1.0".to_string(), status: "pending".to_string(), namespace: "".to_string(), target_date: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 106, "method": "tools/call", "params": { "name": "update_milestone", "arguments": { "id": "v1.0", "status": "completed" } } }); let _ = handler.handle_request(req).await.unwrap(); let ms = handler.state.milestones.read(); assert_eq!(ms[0].status, "completed"); } #[tokio::test] async fn test_handle_list_milestones() { let handler = setup_test_handler("list_milestones"); handler.state.milestones.modify(|ms| { ms.push(crate::models::Milestone { id: "v2.0".to_string(), title: "v2.0".to_string(), status: "pending".to_string(), namespace: "".to_string(), target_date: None, }); }); let req = json!({ "jsonrpc": "2.0", "id": 107, "method": "tools/call", "params": { "name": "list_milestones", "arguments": {} } }); let response = handler.handle_request(req).await.unwrap(); let content = response["result"]["content"][0]["text"].as_str().unwrap(); assert!(content.contains("v2.0")); } #[tokio::test] async fn test_handle_add_pr_checklist_item() { let handler = setup_test_handler("add_pr_checklist_item"); let req = json!({ "jsonrpc": "2.0", "id": 108, "method": "tools/call", "params": { "name": "add_pr_checklist_item", "arguments": { "description": "Check tests", "namespace": "global" } } }); let _ = handler.handle_request(req).await.unwrap(); let cl = handler.state.pr_checklists.read(); assert_eq!(cl.len(), 1); assert_eq!(cl[0].description, "Check tests"); } #[tokio::test] async fn test_handle_clear_pr_checklist() { let handler = setup_test_handler("clear_pr_checklist"); handler.state.pr_checklists.modify(|cl| { cl.push(crate::models::PrChecklistItem { namespace: "global".to_string(), id: "item".to_string(), description: "".to_string(), }); }); let req = json!({ "jsonrpc": "2.0", "id": 109, "method": "tools/call", "params": { "name": "clear_pr_checklist", "arguments": { "namespace": "global" } } }); let _ = handler.handle_request(req).await.unwrap(); let cl = handler.state.pr_checklists.read(); assert!(cl.is_empty()); } #[tokio::test] async fn test_handle_clear_handoff_memos() { let handler = setup_test_handler("clear_handoff_memos"); handler.state.handoff_memos.modify(|ms| { ms.push(crate::models::HandoffMemo { id: "memo".to_string(), namespace: "".to_string(), content: "".to_string(), author: "".to_string(), timestamp: 0, }); }); let req = json!({ "jsonrpc": "2.0", "id": 110, "method": "tools/call", "params": { "name": "clear_handoff_memos", "arguments": { "ids": ["memo"] } } }); let _ = handler.handle_request(req).await.unwrap(); let ms = handler.state.handoff_memos.read(); assert!(ms.is_empty()); } #[tokio::test] async fn test_handle_clear_sticky_notes() { let handler = setup_test_handler("clear_sticky_notes"); handler.state.sticky.modify(|ns| { ns.push(crate::models::StickyNote { content: "".to_string(), timestamp: 0, }); }); let req = json!({ "jsonrpc": "2.0", "id": 111, "method": "tools/call", "params": { "name": "clear_sticky_notes", "arguments": {} } }); let _ = handler.handle_request(req).await.unwrap(); let ns = handler.state.sticky.read(); assert!(ns.is_empty()); } #[tokio::test] async fn test_handle_delete_sticky_note() { let handler = setup_test_handler("delete_sticky_note"); handler.state.sticky.modify(|ns| { ns.push(crate::models::StickyNote { content: "note1".to_string(), timestamp: 0, }); ns.push(crate::models::StickyNote { content: "note2".to_string(), timestamp: 0, }); }); let req = json!({ "jsonrpc": "2.0", "id": 112, "method": "tools/call", "params": { "name": "delete_sticky_note", "arguments": { "index": 1 } } }); let _ = handler.handle_request(req).await.unwrap(); let ns = handler.state.sticky.read(); assert_eq!(ns.len(), 1); assert_eq!(ns[0].content, "note2"); } }