use crate::models::*; use crate::state::MemoryState; use crate::tools::*; use serde::{Deserialize, 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 id = req.get("id").cloned().unwrap_or(serde_json::Value::Null); let method = req.get("method").and_then(|m| m.as_str()).unwrap_or(""); match method { "initialize" => { Some(crate::mcp::success( id, serde_json::json!({ "protocolVersion": "2024-11-05", "capabilities": {}, "serverInfo": { "name": "gemini-mcp-memory", "version": "3.0.0" } }), )) } "notifications/initialized" => { None } "tools/list" => { let tools = vec![ 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::("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::("list_active_tasks", "List all currently active tasks."), 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", "", ), crate::mcp::tool_def::("register_environment", "Register details about a specific deployment environment."), crate::mcp::tool_def::( "get_environment_details", "", ), 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", "", ), 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())); let result: Result = match name { "create_entities" => { let req: CreateEntitiesTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; self.state.write_to_local_delta(|g| { for e_val in req.entities { if let Some(e) = if e_val.is_string() { serde_json::from_str(e_val.as_str().unwrap()).ok() } else { serde_json::from_value(e_val).ok() } { let entity: Entity = e; if !entity.name.is_empty() { if let Ok(idx) = self.state.search_index.read() { let _ = idx.index_entity(&entity); } g.entities.insert(entity.name.clone(), entity); } } } }).await; Ok(vec!["Entities created".to_string()][0].clone()) } "create_relations" => { let req: CreateRelationsTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; self.state.write_to_local_delta(|g| { for r_val in req.relations { if let Some(r) = if r_val.is_string() { serde_json::from_str(r_val.as_str().unwrap()).ok() } else { serde_json::from_value(r_val).ok() } { let relation: Relation = r; if !relation.from.is_empty() && !relation.to.is_empty() { g.relations.push(relation); } } } }).await; Ok(vec!["Relations created".to_string()][0].clone()) } "add_observations" => { let req: AddObservationsTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; #[derive(Deserialize)] struct ObsInput { #[serde(rename = "entityName")] entity_name: String, contents: Vec, } let full = self.state.get_full_graph(); self.state.write_to_local_delta(|g| { for o_val in req.observations { if let Some(o) = if o_val.is_string() { serde_json::from_str::(o_val.as_str().unwrap()).ok() } else { serde_json::from_value(o_val).ok() } && 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); } } }).await; Ok(vec!["Observations added".to_string()][0].clone()) } "delete_entities" => { let req: DeleteEntitiesTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; let to_delete: HashSet<_> = req.entity_names.into_iter().collect(); self.state.apply_sync_write(|master| { for name in &to_delete { master.entities.remove(name); } master.relations.retain(|r| { !to_delete.contains(&r.from) && !to_delete.contains(&r.to) }); }).await; Ok(vec!["Entities deleted".to_string()][0].clone()) } "delete_observations" => { let req: DeleteObservationsTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; #[derive(Deserialize)] struct ObsDel { #[serde(rename = "entityName")] entity_name: String, observations: Vec, } self.state.apply_sync_write(|master| { for d_val in req.deletions { if let Some(d) = if d_val.is_string() { serde_json::from_str::(d_val.as_str().unwrap()).ok() } else { serde_json::from_value(d_val).ok() } && 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)); } } }).await; Ok(vec!["Observations deleted".to_string()][0].clone()) } "delete_relations" => { let req: DeleteRelationsTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; self.state.apply_sync_write(|master| { let mut to_rem = HashSet::new(); for r_val in req.relations { if let Some(r) = if r_val.is_string() { serde_json::from_str::(r_val.as_str().unwrap()).ok() } else { serde_json::from_value(r_val).ok() } { to_rem.insert(format!( "{}|{}|{}|{}", r.from, r.to, r.relation_type, r.namespace )); } } master.relations.retain(|r| { !to_rem.contains(&format!( "{}|{}|{}|{}", r.from, r.to, r.relation_type, r.namespace )) }); }).await; Ok(vec!["Relations deleted".to_string()][0].clone()) } "read_graph" => { let req: ReadGraphTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "search_nodes" => { let req: SearchNodesTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "open_nodes" => { let req: OpenNodesTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "log_code_change" => { let req: LogCodeChangeTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Code change logged".to_string()][0].clone()) } "query_recent_changes" => { let data = serde_json::to_string(&self.state.ledger.read()) .unwrap_or_else(|_| "[]".to_string()); Ok(vec![data.to_string()][0].clone()) } "visualize_graph" => { let req: VisualizeGraphTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![output.to_string()][0].clone()) } "add_sticky_note" => { let req: AddStickyNoteTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; self.state.sticky.modify(|notes| { notes.push(StickyNote { timestamp: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), content: req.content, }); }); Ok(vec!["Sticky note added.".to_string()][0].clone()) } "read_sticky_notes" => { let data = serde_json::to_string(&self.state.sticky.read()) .unwrap_or_else(|_| "[]".to_string()); Ok(vec![data.to_string()][0].clone()) } "condense_entity" => { let req: CondenseEntityTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; self.state.apply_sync_write(|master| { if let Some(e) = master.entities.get_mut(&req.entity_name) { e.observations = req.summarized_observations; } }).await; Ok(vec!["Entity condensed".to_string()][0].clone()) } "add_task" => { let req: AddTaskTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; let now = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(); let id = uuid::Uuid::new_v4().to_string(); let task = Task { id: id.clone(), title: req.title, status: "pending".to_string(), description: req.description, created_at: now, updated_at: now, git_branch: req.git_branch, }; if let Ok(idx) = self.state.search_index.read() { let _ = idx.index_task(&task); } self.state.tasks.modify(|tasks| { tasks.push(task); }); Ok(vec![format!("Task added with ID: {}", id).to_string()][0].clone()) } "update_task_status" => { let req: UpdateTaskStatusTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; let mut found = false; self.state.tasks.modify(|tasks| { for t in tasks.iter_mut() { if t.id == req.id { t.status = req.status.clone(); t.updated_at = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(); found = true; break; } } }); if found { Ok(vec!["Task updated.".to_string()][0].clone()) } else { Ok(vec!["Task not found.".to_string()][0].clone()) } } "list_active_tasks" => { let req: ListActiveTasksTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "store_snippet" => { let req: StoreSnippetTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; self.state.snippets.modify(|snippets| { snippets.retain(|s| s.name != req.name); snippets.push(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(), }); }); Ok(vec![format!("Snippet '{}' stored.", req.name).to_string()][0].clone()) } "search_snippets" => { let req: SearchSnippetsTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "delete_snippet" => { let req: DeleteSnippetTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Snippet deleted.".to_string()][0].clone()) } else { Ok(vec!["Snippet not found.".to_string()][0].clone()) } } "log_decision" => { let req: LogDecisionTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; let mut id = String::new(); self.state.adrs.modify(|adrs| { id = format!("ADR-{:04}", adrs.len() + 1); adrs.push(Adr { id: id.clone(), title: req.title, context: req.context, decision: req.decision, consequence: req.consequence, timestamp: SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(), }); }); Ok(vec![format!("Decision logged as {}", id).to_string()][0].clone()) } "query_decisions" => { let req: QueryDecisionsTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "merge_entities" => { let req: MergeEntitiesTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; self.state.apply_sync_write(|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()); }).await; Ok(vec!["Entities merged".to_string()][0].clone()) } "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(vec![data.to_string()][0].clone()) } "learn_preference" => { let req: LearnPreferenceTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Preference learned".to_string()][0].clone()) } "read_preferences" => { let prefs = self.state.prefs.read(); let data = serde_json::to_string(&prefs).unwrap_or_default(); Ok(vec![data.to_string()][0].clone()) } "log_error_fix" => { let req: LogErrorFixTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Error fix logged".to_string()][0].clone()) } "search_error_fixes" => { let req: SearchErrorFixesTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "pin_file" => { let req: PinFileTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["File pinned".to_string()][0].clone()) } "unpin_file" => { let req: UnpinFileTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; self.state.pinned_files.modify(|pinned| { pinned.retain(|p| { !(p.namespace == req.namespace && p.file_path == req.file_path) }) }); Ok(vec!["File unpinned".to_string()][0].clone()) } "list_pinned_files" => { let req: ListPinnedFilesTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "add_session_summary" => { let req: AddSessionSummaryTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Session summary added".to_string()][0].clone()) } "get_project_timeline" => { let req: GetProjectTimelineTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "leave_handoff_memo" => { let req: LeaveHandoffMemoTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Handoff memo left".to_string()][0].clone()) } "read_handoff_memos" => { let req: ReadHandoffMemosTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "clear_handoff_memos" => { let req: ClearHandoffMemosTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; let ids: HashSet<_> = req.ids.into_iter().collect(); self.state .handoff_memos .modify(|memos| memos.retain(|m| !ids.contains(&m.id))); Ok(vec!["Handoff memos cleared".to_string()][0].clone()) } "update_env_fingerprint" => { let req: UpdateEnvFingerprintTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Env fingerprint updated".to_string()][0].clone()) } "read_env_fingerprint" => { let req: ReadEnvFingerprintTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } else { Ok(vec!["{}".to_string()][0].clone()) } } "log_env_requirement" => { let req: LogEnvRequirementTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Env requirement logged".to_string()][0].clone()) } "add_milestone" => { let req: AddMilestoneTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Milestone added".to_string()][0].clone()) } "update_milestone" => { let req: UpdateMilestoneTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Milestone updated".to_string()][0].clone()) } else { Ok(vec!["Milestone not found".to_string()][0].clone()) } } "list_milestones" => { let req: ListMilestonesTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "generate_standup_report" => { let req: GenerateStandupReportTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![report.to_string().to_string()][0].clone()) } "register_environment" => { let req: RegisterEnvironmentTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Environment registered".to_string()][0].clone()) } "get_environment_details" => { let req: GetEnvironmentDetailsTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "add_pr_checklist_item" => { let req: AddPrChecklistItemTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["PR checklist item added".to_string()][0].clone()) } "get_pr_checklist" => { let req: GetPrChecklistTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "clear_pr_checklist" => { let req: ClearPrChecklistTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; self.state .pr_checklists .modify(|items| items.retain(|i| i.namespace != req.namespace)); Ok(vec!["PR checklist cleared".to_string()][0].clone()) } "log_tech_debt" => { let req: LogTechDebtTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Tech debt logged".to_string()][0].clone()) } "resolve_tech_debt" => { let req: ResolveTechDebtTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Tech debt resolved".to_string()][0].clone()) } else { Ok(vec!["Tech debt not found".to_string()][0].clone()) } } "list_tech_debt" => { let req: ListTechDebtTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "save_context_workspace" => { let req: SaveContextWorkspaceTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec!["Context workspace saved".to_string()][0].clone()) } "load_context_workspace" => { let req: LoadContextWorkspaceTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "list_context_workspaces" => { let req: ListContextWorkspacesTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![data.to_string()][0].clone()) } "omni_search" => { let req: OmniSearchTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![report.to_string().to_string()][0].clone()) } "get_project_health" => { let req: GetProjectHealthTool = match parse_args(args.clone()) { Ok(r) => r, Err(e) => { return Some(crate::mcp::success( id.clone(), serde_json::json!({"isError": true, "content": [{"type": "text", "text": e}] }), )); } }; 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(vec![report.to_string().to_string()][0].clone()) } _ => 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 { return Some(crate::mcp::error(id, -32601, "Method not found")); } None } } } }