diff --git a/Cargo.lock b/Cargo.lock index 4fc7ca8..d71def1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -157,6 +157,15 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + [[package]] name = "bitflags" version = "2.13.1" @@ -392,6 +401,20 @@ dependencies = [ "syn 3.0.5", ] +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "datasketches" version = "0.2.0" @@ -626,6 +649,12 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + [[package]] name = "hashbrown" version = "0.16.1" @@ -981,7 +1010,7 @@ version = "0.16.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f66e8d5d03f609abc3a39e6f08e4164ebf1447a732906d39eb9b99b7919ef39" dependencies = [ - "hashbrown", + "hashbrown 0.16.1", ] [[package]] @@ -1008,10 +1037,13 @@ version = "0.1.0" dependencies = [ "async-trait", "axum", + "bincode", "clap", + "dashmap", "dirs", "futures-util", "glob", + "redb", "reqwest", "schemars", "serde", @@ -1357,6 +1389,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "redb" +version = "4.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de6c3b63e007e90ce536ec2ae4690826136a20ec8dbbbb400daef1bb999d2e36" +dependencies = [ + "libc", +] + [[package]] name = "redox_syscall" version = "0.5.18" diff --git a/server/Cargo.toml b/server/Cargo.toml index 0e94f20..dd6422f 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -6,10 +6,13 @@ edition = "2024" [dependencies] async-trait = "0.1.92" axum = "0.8" +bincode = "1.3.3" clap = { version = "4.6.6", features = ["derive"] } +dashmap = "6.2.1" dirs = "6.0.0" futures-util = "0.3.34" glob = "0.3.4" +redb = "4.2.0" reqwest = { version = "0.12", default-features = false, features = ["stream", "rustls-tls"] } schemars = "1.2.2" serde = { version = "1.0.229", features = ["derive"] } diff --git a/server/src/dashboard.html b/server/src/dashboard.html index ca17e42..3d3d503 100644 --- a/server/src/dashboard.html +++ b/server/src/dashboard.html @@ -5,13 +5,13 @@ diff --git a/server/src/handlers.rs b/server/src/handlers.rs index 2f2da52..e74b781 100644 --- a/server/src/handlers.rs +++ b/server/src/handlers.rs @@ -2,19 +2,15 @@ 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 { +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, } @@ -23,1168 +19,1568 @@ 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" => { - return Some(crate::mcp::success(id, serde_json::json!({ - "protocolVersion": "2024-11-05", - "capabilities": {}, - "serverInfo": { - "name": "gemini-mcp-memory", - "version": "3.0.0" - } - }))); + Some(crate::mcp::success( + id, + serde_json::json!({ + "protocolVersion": "2024-11-05", + "capabilities": {}, + "serverInfo": { + "name": "gemini-mcp-memory", + "version": "3.0.0" + } + }), + )) } "notifications/initialized" => { - return None; + None } "tools/list" => { - let tools = vec![ crate::mcp::tool_def::("create_entities", ""), crate::mcp::tool_def::("create_relations", ""), crate::mcp::tool_def::("add_observations", ""), crate::mcp::tool_def::("delete_entities", ""), crate::mcp::tool_def::("delete_observations", ""), crate::mcp::tool_def::("delete_relations", ""), crate::mcp::tool_def::("read_graph", ""), crate::mcp::tool_def::("search_nodes", ""), crate::mcp::tool_def::("open_nodes", ""), crate::mcp::tool_def::("log_code_change", ""), crate::mcp::tool_def::("query_recent_changes", ""), crate::mcp::tool_def::("visualize_graph", ""), crate::mcp::tool_def::("add_sticky_note", ""), crate::mcp::tool_def::("read_sticky_notes", ""), crate::mcp::tool_def::("condense_entity", ""), crate::mcp::tool_def::("add_task", ""), crate::mcp::tool_def::("update_task_status", ""), crate::mcp::tool_def::("list_active_tasks", ""), crate::mcp::tool_def::("store_snippet", ""), crate::mcp::tool_def::("search_snippets", ""), crate::mcp::tool_def::("delete_snippet", ""), crate::mcp::tool_def::("log_decision", ""), crate::mcp::tool_def::("query_decisions", ""), crate::mcp::tool_def::("merge_entities", ""), crate::mcp::tool_def::("find_orphans", ""), crate::mcp::tool_def::("learn_preference", ""), crate::mcp::tool_def::("read_preferences", ""), crate::mcp::tool_def::("log_error_fix", ""), crate::mcp::tool_def::("search_error_fixes", ""), crate::mcp::tool_def::("pin_file", ""), crate::mcp::tool_def::("unpin_file", ""), crate::mcp::tool_def::("list_pinned_files", ""), crate::mcp::tool_def::("add_session_summary", ""), crate::mcp::tool_def::("get_project_timeline", ""), crate::mcp::tool_def::("leave_handoff_memo", ""), crate::mcp::tool_def::("read_handoff_memos", ""), crate::mcp::tool_def::("clear_handoff_memos", ""), crate::mcp::tool_def::("update_env_fingerprint", ""), crate::mcp::tool_def::("read_env_fingerprint", ""), crate::mcp::tool_def::("log_env_requirement", ""), crate::mcp::tool_def::("add_milestone", ""), crate::mcp::tool_def::("update_milestone", ""), crate::mcp::tool_def::("list_milestones", ""), crate::mcp::tool_def::("generate_standup_report", ""), crate::mcp::tool_def::("register_environment", ""), crate::mcp::tool_def::("get_environment_details", ""), crate::mcp::tool_def::("add_pr_checklist_item", ""), crate::mcp::tool_def::("get_pr_checklist", ""), crate::mcp::tool_def::("clear_pr_checklist", ""), crate::mcp::tool_def::("log_tech_debt", ""), crate::mcp::tool_def::("resolve_tech_debt", ""), crate::mcp::tool_def::("list_tech_debt", ""), crate::mcp::tool_def::("save_context_workspace", ""), crate::mcp::tool_def::("load_context_workspace", ""), crate::mcp::tool_def::("list_context_workspaces", ""), crate::mcp::tool_def::("omni_search", ""), crate::mcp::tool_def::("get_project_health", ""), + let tools = vec![ + crate::mcp::tool_def::("create_entities", ""), + crate::mcp::tool_def::("create_relations", ""), + crate::mcp::tool_def::("add_observations", ""), + crate::mcp::tool_def::("delete_entities", ""), + crate::mcp::tool_def::("delete_observations", ""), + crate::mcp::tool_def::("delete_relations", ""), + crate::mcp::tool_def::("read_graph", ""), + crate::mcp::tool_def::("search_nodes", ""), + crate::mcp::tool_def::("open_nodes", ""), + crate::mcp::tool_def::("log_code_change", ""), + crate::mcp::tool_def::("query_recent_changes", ""), + crate::mcp::tool_def::("visualize_graph", ""), + crate::mcp::tool_def::("add_sticky_note", ""), + crate::mcp::tool_def::("read_sticky_notes", ""), + crate::mcp::tool_def::("condense_entity", ""), + crate::mcp::tool_def::("add_task", ""), + crate::mcp::tool_def::("update_task_status", ""), + crate::mcp::tool_def::("list_active_tasks", ""), + crate::mcp::tool_def::("store_snippet", ""), + crate::mcp::tool_def::("search_snippets", ""), + crate::mcp::tool_def::("delete_snippet", ""), + crate::mcp::tool_def::("log_decision", ""), + crate::mcp::tool_def::("query_decisions", ""), + crate::mcp::tool_def::("merge_entities", ""), + crate::mcp::tool_def::("find_orphans", ""), + crate::mcp::tool_def::("learn_preference", ""), + crate::mcp::tool_def::("read_preferences", ""), + crate::mcp::tool_def::("log_error_fix", ""), + crate::mcp::tool_def::("search_error_fixes", ""), + crate::mcp::tool_def::("pin_file", ""), + crate::mcp::tool_def::("unpin_file", ""), + crate::mcp::tool_def::("list_pinned_files", ""), + crate::mcp::tool_def::("add_session_summary", ""), + crate::mcp::tool_def::("get_project_timeline", ""), + crate::mcp::tool_def::("leave_handoff_memo", ""), + crate::mcp::tool_def::("read_handoff_memos", ""), + crate::mcp::tool_def::("clear_handoff_memos", ""), + crate::mcp::tool_def::("update_env_fingerprint", ""), + crate::mcp::tool_def::("read_env_fingerprint", ""), + crate::mcp::tool_def::("log_env_requirement", ""), + crate::mcp::tool_def::("add_milestone", ""), + crate::mcp::tool_def::("update_milestone", ""), + crate::mcp::tool_def::("list_milestones", ""), + crate::mcp::tool_def::( + "generate_standup_report", + "", + ), + crate::mcp::tool_def::("register_environment", ""), + crate::mcp::tool_def::( + "get_environment_details", + "", + ), + crate::mcp::tool_def::("add_pr_checklist_item", ""), + crate::mcp::tool_def::("get_pr_checklist", ""), + crate::mcp::tool_def::("clear_pr_checklist", ""), + crate::mcp::tool_def::("log_tech_debt", ""), + crate::mcp::tool_def::("resolve_tech_debt", ""), + crate::mcp::tool_def::("list_tech_debt", ""), + crate::mcp::tool_def::("save_context_workspace", ""), + crate::mcp::tool_def::("load_context_workspace", ""), + crate::mcp::tool_def::( + "list_context_workspaces", + "", + ), + crate::mcp::tool_def::("omni_search", ""), + crate::mcp::tool_def::("get_project_health", ""), ]; - return Some(crate::mcp::success(id, serde_json::json!({ "tools": tools }))); + 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 { + let args = params + .get("arguments") + .cloned() + .unwrap_or(serde_json::Value::Object(Default::default())); - "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); + 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); + } } - g.entities.insert(entity.name.clone(), entity); + } + }); + 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); + } + } + } + }); + 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); + } + } + }); + 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) + }); + }); + 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)); + } + } + }); + 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 + )) + }); + }); + 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()); } } - } - }); - 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); + 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()) } - }); - 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); + "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()); + } } - } - }); - 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)); - }); - 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)); + 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); + } } - } - }); - 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 + 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 )); } - } - master.relations.retain(|r| { - !to_rem.contains(&format!( - "{}|{}|{}|{}", - r.from, r.to, r.relation_type, r.namespace - )) - }); - }); - 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" { - if let Some(e) = full.entities.get(&id) { - result.entities.insert(id, e.clone()); + 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 + )); } - } - } - 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; - } - }); - 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 output == "graph TD;\n" { + output = "No nodes found to visualize.".to_string(); } + Ok(vec![output.to_string()][0].clone()) } - }); - 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() + "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; + } + }); + 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(), - }); - }); - 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); + .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()) } - } - 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()); + "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 { - let mut new_tgt = src.clone(); - new_tgt.name = req.target_entity.clone(); - master.entities.insert(req.target_entity.clone(), new_tgt); + Ok(vec!["Task not found.".to_string()][0].clone()) } } - for r in &mut master.relations { - if r.from == req.source_entity { - r.from = req.target_entity.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()) + }); } - if r.to == req.source_entity { - r.to = req.target_entity.clone(); + 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()) } } - master.relations = MemoryState::unique_items(master.relations.clone()); - }); - 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; + "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()); + }); + 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()) } } - }); - 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; + "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()) } } - }); - 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() + "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(), - }); - }); - 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![] - }; + .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" { - if let Some(e) = full.entities.get(id) { - kg.entities.insert(id.clone(), e.clone()); + 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); + } } - } - } - 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 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!({ + "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()) + } - 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) => return Some(crate::mcp::success(id, serde_json::json!({ - "content": [{ "type": "text", "text": text }] - }))), - Err(e) => return Some(crate::mcp::success(id, serde_json::json!({ - "isError": true, - "content": [{ "type": "text", "text": e }] - }))) + 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")); } - return None; + None } } } diff --git a/server/src/main.rs b/server/src/main.rs index 9955d96..b1f1ec2 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -1,10 +1,10 @@ mod handlers; -mod models; mod mcp; +mod models; +mod search; mod state; mod store; mod tools; -mod search; use crate::handlers::MemoryHandler; use crate::models::*; @@ -29,6 +29,10 @@ struct Cli { target: Option, #[arg(long)] daemon: bool, + #[arg(long)] + exit: bool, + #[arg(long)] + restart: bool, } #[derive(Subcommand)] @@ -71,7 +75,6 @@ enum GateCommands { }, } - async fn reconcile_worker(state: Arc) { loop { sleep(Duration::from_secs(5)).await; @@ -104,21 +107,17 @@ async fn reconcile_worker(state: Arc) { } } - - - - use axum::{ - extract::{State, Query}, + Json, Router, + extract::{Query, State}, response::sse::{Event, Sse}, routing::{get, post}, - Json, Router, }; use futures_util::stream::Stream; use std::convert::Infallible; +use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; -use std::sync::atomic::{AtomicUsize, Ordering}; struct AppState { handler: Arc, @@ -140,58 +139,72 @@ fn run_server(state: Arc) -> Result<(), Box> .route("/sse", get(sse_handler)) .route("/messages", post(message_handler)) .route("/health", get(health_handler)) - .route("/", get(|| async move { - axum::response::Html(include_str!("dashboard.html")) - })) - .route("/api/stats", get({ - let state_clone = app_state.handler.state.clone(); - move || async move { - let (entities, relations) = { - let graph = state_clone.get_full_graph(); - (graph.entities.len(), graph.relations.len()) - }; - let tasks = state_clone.tasks.read().len(); - let snippets = state_clone.snippets.read().len(); - let tech_debts = state_clone.tech_debts.read().len(); - let adrs = state_clone.adrs.read().len(); - - let ledger = state_clone.ledger.read().len(); - let sticky = state_clone.sticky.read().len(); - let error_fixes = state_clone.error_fixes.read().len(); - let pinned_files = state_clone.pinned_files.read().len(); - let session_summaries = state_clone.session_summaries.read().len(); - let handoff_memos = state_clone.handoff_memos.read().len(); - let env_fingerprints = state_clone.env_fingerprints.read().len(); - let env_requirements = state_clone.env_requirements.read().len(); - let milestones = state_clone.milestones.read().len(); - let environments = state_clone.environments.read().len(); - let pr_checklists = state_clone.pr_checklists.read().len(); - let gates = state_clone.gates.read().len(); - let context_workspaces = state_clone.context_workspaces.read().len(); + .route( + "/shutdown", + post(|| async move { + std::thread::spawn(|| { + std::thread::sleep(std::time::Duration::from_millis(100)); + std::process::exit(0); + }); + "Shutting down..." + }), + ) + .route( + "/", + get(|| async move { axum::response::Html(include_str!("dashboard.html")) }), + ) + .route( + "/api/stats", + get({ + let state_clone = app_state.handler.state.clone(); + move || async move { + let (entities, relations) = { + let graph = state_clone.get_full_graph(); + (graph.entities.len(), graph.relations.len()) + }; + let tasks = state_clone.tasks.read().len(); + let snippets = state_clone.snippets.read().len(); + let tech_debts = state_clone.tech_debts.read().len(); + let adrs = state_clone.adrs.read().len(); - axum::Json(serde_json::json!({ - "entities": entities, - "relations": relations, - "tasks": tasks, - "snippets": snippets, - "tech_debts": tech_debts, - "adrs": adrs, - "ledger": ledger, - "sticky": sticky, - "error_fixes": error_fixes, - "pinned_files": pinned_files, - "session_summaries": session_summaries, - "handoff_memos": handoff_memos, - "env_fingerprints": env_fingerprints, - "env_requirements": env_requirements, - "milestones": milestones, - "environments": environments, - "pr_checklists": pr_checklists, - "gates": gates, - "context_workspaces": context_workspaces - })) - } - })) + let ledger = state_clone.ledger.read().len(); + let sticky = state_clone.sticky.read().len(); + let error_fixes = state_clone.error_fixes.read().len(); + let pinned_files = state_clone.pinned_files.read().len(); + let session_summaries = state_clone.session_summaries.read().len(); + let handoff_memos = state_clone.handoff_memos.read().len(); + let env_fingerprints = state_clone.env_fingerprints.read().len(); + let env_requirements = state_clone.env_requirements.read().len(); + let milestones = state_clone.milestones.read().len(); + let environments = state_clone.environments.read().len(); + let pr_checklists = state_clone.pr_checklists.read().len(); + let gates = state_clone.gates.read().len(); + let context_workspaces = state_clone.context_workspaces.read().len(); + + axum::Json(serde_json::json!({ + "entities": entities, + "relations": relations, + "tasks": tasks, + "snippets": snippets, + "tech_debts": tech_debts, + "adrs": adrs, + "ledger": ledger, + "sticky": sticky, + "error_fixes": error_fixes, + "pinned_files": pinned_files, + "session_summaries": session_summaries, + "handoff_memos": handoff_memos, + "env_fingerprints": env_fingerprints, + "env_requirements": env_requirements, + "milestones": milestones, + "environments": environments, + "pr_checklists": pr_checklists, + "gates": gates, + "context_workspaces": context_workspaces + })) + } + }), + ) .with_state(app_state); let mut retries = 0; @@ -199,16 +212,48 @@ fn run_server(state: Arc) -> Result<(), Box> match tokio::net::TcpListener::bind("127.0.0.1:3000").await { Ok(l) => break l, Err(e) => { + // Check if it's already running and healthy + if let Ok(mut stream) = std::net::TcpStream::connect("127.0.0.1:3000") { + use std::io::{Read, Write}; + let _ = stream.write_all( + b"GET /health HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n", + ); + let mut response = String::new(); + let _ = stream.read_to_string(&mut response); + if response.contains("200 OK") { + // Already healthy! Just exit cleanly instead of panicking/retrying loop. + std::process::exit(0); + } + } + retries += 1; if retries > 15 { - let log_path = dirs::home_dir().unwrap_or_default().join(".gemini/mcp_memory/daemon_fatal.log"); - let _ = std::fs::write(&log_path, format!("FATAL: Could not bind to 127.0.0.1:3000 after 15 seconds: {}\n", e)); + let log_path = dirs::home_dir() + .unwrap_or_default() + .join(".gemini/mcp_memory/daemon_fatal.log"); + let _ = std::fs::write( + &log_path, + format!( + "FATAL: Could not bind to 127.0.0.1:3000 after 15 seconds: {}\n", + e + ), + ); std::process::exit(1); } - let log_path = dirs::home_dir().unwrap_or_default().join(".gemini/mcp_memory/daemon_error.log"); - if let Ok(mut file) = std::fs::OpenOptions::new().create(true).append(true).open(&log_path) { + let log_path = dirs::home_dir() + .unwrap_or_default() + .join(".gemini/mcp_memory/daemon_error.log"); + if let Ok(mut file) = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&log_path) + { use std::io::Write; - let _ = writeln!(file, "Failed to bind to 127.0.0.1:3000 (attempt {}): {}. Retrying in 1s...", retries, e); + let _ = writeln!( + file, + "Failed to bind to 127.0.0.1:3000 (attempt {}): {}. Retrying in 1s...", + retries, e + ); } tokio::time::sleep(std::time::Duration::from_secs(1)).await; } @@ -216,7 +261,9 @@ fn run_server(state: Arc) -> Result<(), Box> }; eprintln!("MCP Memory Server running on http://127.0.0.1:3000/sse"); if let Err(e) = axum::serve(listener, app).await { - let log_path = dirs::home_dir().unwrap_or_default().join(".gemini/mcp_memory/daemon_error.log"); + let log_path = dirs::home_dir() + .unwrap_or_default() + .join(".gemini/mcp_memory/daemon_error.log"); let _ = std::fs::write(&log_path, format!("Server crashed: {}\n", e)); } Ok(()) @@ -228,11 +275,19 @@ async fn sse_handler( ) -> Sse>> { let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst)); let (tx, rx) = mpsc::channel::>(100); - - state.clients.write().unwrap().insert(session_id.clone(), tx.clone()); - - let _ = tx.send(Ok(Event::default().event("endpoint").data(format!("/messages?sessionId={}", session_id)))).await; - + + state + .clients + .write() + .unwrap() + .insert(session_id.clone(), tx.clone()); + + let _ = tx + .send(Ok(Event::default() + .event("endpoint") + .data(format!("/messages?sessionId={}", session_id)))) + .await; + let stream = ReceiverStream::new(rx); Sse::new(stream).keep_alive(axum::response::sse::KeepAlive::new()) } @@ -255,32 +310,56 @@ async fn message_handler( let handler = Arc::clone(&state.handler); let session_id = query.session_id.clone(); let clients = Arc::clone(&state); - + tokio::spawn(async move { if let Some(response) = handler.handle_request(payload).await { let tx_opt = clients.clients.read().unwrap().get(&session_id).cloned(); if let Some(tx) = tx_opt { let data = serde_json::to_string(&response).unwrap(); - let _ = tx.send(Ok(Event::default().event("message").data(data))).await; + let _ = tx + .send(Ok(Event::default().event("message").data(data))) + .await; } } }); - + axum::http::StatusCode::ACCEPTED } mod proxy; - fn main() -> Result<(), Box> { let cli = Cli::parse(); + if cli.exit { + if let Ok(mut stream) = std::net::TcpStream::connect("127.0.0.1:3000") { + use std::io::Write; + let _ = stream.write_all( + b"POST /shutdown HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n", + ); + } + println!("Sent shutdown request to server."); + return Ok(()); + } + + if cli.restart { + if let Ok(mut stream) = std::net::TcpStream::connect("127.0.0.1:3000") { + use std::io::Write; + let _ = stream.write_all( + b"POST /shutdown HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n", + ); + println!("Sent shutdown request to existing server. Waiting for it to exit..."); + std::thread::sleep(std::time::Duration::from_millis(1500)); + } + return Ok(()); + } + #[cfg(target_os = "windows")] { use std::os::windows::process::CommandExt; if !cli.daemon { loop { - if let Err(_) = std::net::TcpListener::bind("127.0.0.1:3000") { + if std::net::TcpListener::bind("127.0.0.1:3000").is_err() { // Port in use, become a stub proxy! let target_url = cli.target.as_deref().unwrap_or("http://127.0.0.1:3000"); match proxy::run_proxy(target_url) { @@ -293,7 +372,8 @@ fn main() -> Result<(), Box> { } } else { // Port is free. We must spawn the daemon, then loop again to become proxy - std::process::Command::new(std::env::current_exe().unwrap()) + #[allow(clippy::zombie_processes)] + let _ = std::process::Command::new(std::env::current_exe().unwrap()) .arg("--daemon") .stdin(std::process::Stdio::null()) .stdout(std::process::Stdio::null()) @@ -332,7 +412,7 @@ fn main() -> Result<(), Box> { session_graph: RwLock::new(KnowledgeGraph::default()), base_dir: base.clone(), master_cache: RwLock::new((KnowledgeGraph::default(), SystemTime::UNIX_EPOCH)), - search_index: RwLock::new(crate::search::MemoryIndex::new().unwrap()), + search_index: RwLock::new(crate::search::MemoryIndex::new(&base).unwrap()), ledger: Store::new(base.join("audit_ledger.json")), sticky: Store::new(base.join("sticky_notes.json")), tasks: Store::new(base.join("tasks.json")), @@ -357,75 +437,101 @@ fn main() -> Result<(), Box> { if let Some(command) = cli.command { match command { - Commands::Gate { subcmd } => { - match subcmd { - GateCommands::Set { action, target, namespace, params, authorize, block, reason } => { - let status = if authorize { "authorized".to_string() } else if block { "blocked".to_string() } else { "pending".to_string() }; - let mut param_map = HashMap::new(); - for p in params { - if let Some((k, v)) = p.split_once('=') { - param_map.insert(k.to_string(), v.to_string()); - } + Commands::Gate { subcmd } => match subcmd { + GateCommands::Set { + action, + target, + namespace, + params, + authorize, + block, + reason, + } => { + let status = if authorize { + "authorized".to_string() + } else if block { + "blocked".to_string() + } else { + "pending".to_string() + }; + let mut param_map = HashMap::new(); + for p in params { + if let Some((k, v)) = p.split_once('=') { + param_map.insert(k.to_string(), v.to_string()); } - let record = GateRecord { - id: uuid::Uuid::new_v4().to_string(), - action: action.clone(), - target: target.clone(), - namespace, - params: param_map, - status, - reason, - timestamp: SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(), - }; - state.gates.modify(|gates| { - gates.retain(|g| !(g.action == record.action && g.target == record.target)); - gates.push(record); - }); - println!("Gate state updated."); - std::process::exit(0); } - GateCommands::Verify { action, target, namespace, params, consume } => { - let mut param_map = HashMap::new(); - for p in params { - if let Some((k, v)) = p.split_once('=') { - param_map.insert(k.to_string(), v.to_string()); + let record = GateRecord { + id: uuid::Uuid::new_v4().to_string(), + action: action.clone(), + target: target.clone(), + namespace, + params: param_map, + status, + reason, + timestamp: SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(), + }; + state.gates.modify(|gates| { + gates.retain(|g| !(g.action == record.action && g.target == record.target)); + gates.push(record); + }); + println!("Gate state updated."); + std::process::exit(0); + } + GateCommands::Verify { + action, + target, + namespace, + params, + consume, + } => { + let mut param_map = HashMap::new(); + for p in params { + if let Some((k, v)) = p.split_once('=') { + param_map.insert(k.to_string(), v.to_string()); + } + } + let mut found = None; + let mut to_remove = None; + state.gates.modify(|gates| { + if let Some(idx) = gates.iter().position(|g| { + g.action == action + && g.target == target + && g.namespace == namespace + && g.params == param_map + }) { + found = Some(gates[idx].clone()); + if consume { + to_remove = Some(idx); } } - let mut found = None; - let mut to_remove = None; - state.gates.modify(|gates| { - if let Some(idx) = gates.iter().position(|g| g.action == action && g.target == target && g.namespace == namespace && g.params == param_map) { - found = Some(gates[idx].clone()); - if consume { - to_remove = Some(idx); - } - } - if let Some(idx) = to_remove { - gates.remove(idx); - } - }); - - match found { - Some(record) => { - if record.status == "authorized" { - std::process::exit(0); + if let Some(idx) = to_remove { + gates.remove(idx); + } + }); + + match found { + Some(record) => { + if record.status == "authorized" { + std::process::exit(0); + } else { + if let Some(r) = record.reason { + eprintln!("❌ Action blocked. Reason: {}", r); } else { - if let Some(r) = record.reason { - eprintln!("❌ Action blocked. Reason: {}", r); - } else { - eprintln!("❌ Action blocked."); - } - std::process::exit(1); + eprintln!("❌ Action blocked."); } + std::process::exit(1); } - None => { - eprintln!("❌ Action not yet authorized (no gate record found)."); - std::process::exit(2); - } + } + None => { + eprintln!("❌ Action not yet authorized (no gate record found)."); + std::process::exit(2); } } } - } + }, } } diff --git a/server/src/models.rs b/server/src/models.rs index b7509af..04dacff 100644 --- a/server/src/models.rs +++ b/server/src/models.rs @@ -165,8 +165,7 @@ pub struct ContextWorkspace { pub saved_at: u64, } -#[derive(Debug, Clone, Serialize, Deserialize)] -#[derive(Default)] +#[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct GateRecord { pub id: String, pub action: String, diff --git a/server/src/proxy.rs b/server/src/proxy.rs index 237bafa..e366780 100644 --- a/server/src/proxy.rs +++ b/server/src/proxy.rs @@ -1,9 +1,9 @@ -use tokio_util::io::StreamReader; -use tokio::io::AsyncBufReadExt; use futures_util::StreamExt; use std::sync::Arc; +use tokio::io::AsyncBufReadExt; use tokio::sync::RwLock; use tokio::sync::mpsc; +use tokio_util::io::StreamReader; pub fn run_proxy(target_url: &str) -> Result> { let rt = tokio::runtime::Runtime::new().unwrap(); @@ -16,7 +16,9 @@ pub fn run_proxy(target_url: &str) -> Result Result Result Result { - if resp.status() == reqwest::StatusCode::GONE { - return Ok(true); - } - - let stream = resp.bytes_stream().map(|res| res.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))); - let mut reader = tokio::io::BufReader::new(StreamReader::new(stream)); - let mut line = String::new(); - let mut is_message = false; - let mut is_endpoint = false; - - loop { - tokio::select! { - _ = shutdown_rx.recv() => { - return Ok(false); - } - res = reader.read_line(&mut line) => { - match res { - Ok(bytes) => { - if bytes == 0 { break; } - let trimmed = line.trim(); - if trimmed.starts_with("event: message") { - is_message = true; - is_endpoint = false; - } else if trimmed.starts_with("event: endpoint") { - is_endpoint = true; + let sse_url = format!("{}/sse", target_url); + let client = reqwest::Client::builder().build().unwrap(); + + match client + .get(&sse_url) + .header("Accept", "text/event-stream") + .send() + .await + { + Ok(resp) => { + if resp.status() == reqwest::StatusCode::GONE { + return Ok(true); + } + + let stream = resp.bytes_stream().map(|res| { + res.map_err(std::io::Error::other) + }); + let mut reader = tokio::io::BufReader::new(StreamReader::new(stream)); + let mut line = String::new(); + let mut is_message = false; + let mut is_endpoint = false; + + loop { + tokio::select! { + _ = shutdown_rx.recv() => { + return Ok(false); + } + res = reader.read_line(&mut line) => { + match res { + Ok(bytes) => { + if bytes == 0 { break; } + let trimmed = line.trim(); + if trimmed.starts_with("event: message") { + is_message = true; + is_endpoint = false; + } else if trimmed.starts_with("event: endpoint") { + is_endpoint = true; + is_message = false; + } else if let Some(stripped) = trimmed.strip_prefix("data: ") { + if is_message { + println!("{}", stripped); is_message = false; - } else if trimmed.starts_with("data: ") { - if is_message { - println!("{}", &trimmed[6..]); - is_message = false; - } else if is_endpoint { - let ep = &trimmed[6..]; - let mut p = post_url.write().await; - *p = format!("{}{}", target_url, ep); - is_endpoint = false; - } + } else if is_endpoint { + let mut p = post_url.write().await; + *p = format!("{}{}", target_url, stripped); + is_endpoint = false; } - line.clear(); } - Err(_) => break, + line.clear(); } + Err(_) => break, } } } - *post_url.write().await = String::new(); - return Ok(true); - } - Err(_) => { - return Ok(true); } + *post_url.write().await = String::new(); + Ok(true) } + Err(_) => Ok(true), } }) } diff --git a/server/src/search.rs b/server/src/search.rs index bb02a35..656714a 100644 --- a/server/src/search.rs +++ b/server/src/search.rs @@ -1,13 +1,13 @@ -use tantivy::schema::*; -use tantivy::{doc, Index, IndexWriter, IndexReader, ReloadPolicy}; +use crate::models::{Adr, Entity, Snippet, Task}; use std::sync::Mutex; -use crate::models::{Entity, Task, Adr, Snippet}; +use tantivy::schema::*; +use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, doc}; pub struct MemoryIndex { index: Index, reader: IndexReader, writer: Mutex, - + // Schema fields pub id_field: Field, pub title_field: Field, @@ -17,7 +17,7 @@ pub struct MemoryIndex { } impl MemoryIndex { - pub fn new() -> tantivy::Result { + pub fn new(store_dir: &std::path::Path) -> tantivy::Result { let mut schema_builder = Schema::builder(); let id_field = schema_builder.add_text_field("id", STRING | STORED); let title_field = schema_builder.add_text_field("title", TEXT | STORED); @@ -25,14 +25,17 @@ impl MemoryIndex { let type_field = schema_builder.add_text_field("type", STRING | STORED); let namespace_field = schema_builder.add_text_field("namespace", STRING | STORED); let schema = schema_builder.build(); - - let index = Index::create_in_ram(schema.clone()); + + let index_dir = store_dir.join("tantivy_index"); + std::fs::create_dir_all(&index_dir).unwrap(); + let index = Index::open_in_dir(&index_dir).unwrap_or_else(|_| Index::create_in_dir(&index_dir, schema.clone()).unwrap()); + let writer = index.writer(50_000_000)?; let reader = index .reader_builder() .reload_policy(ReloadPolicy::OnCommitWithDelay) .try_into()?; - + Ok(Self { index, reader, @@ -71,23 +74,43 @@ impl MemoryIndex { Ok(()) } - pub fn search(&self, query: &str, namespace: Option<&str>) -> tantivy::Result> { + pub fn search( + &self, + query: &str, + namespace: Option<&str>, + ) -> tantivy::Result> { let searcher = self.reader.searcher(); - let query_parser = tantivy::query::QueryParser::for_index(&self.index, vec![self.title_field, self.body_field]); + let query_parser = tantivy::query::QueryParser::for_index( + &self.index, + vec![self.title_field, self.body_field], + ); let q = query_parser.parse_query(query)?; - - let top_docs = searcher.search(&q, &tantivy::collector::TopDocs::with_limit(50).order_by_score())?; + + let top_docs = searcher.search( + &q, + &tantivy::collector::TopDocs::with_limit(50).order_by_score(), + )?; let mut results = Vec::new(); for (_score, doc_address) in top_docs { let retrieved_doc = searcher.doc::(doc_address)?; - let id = retrieved_doc.get_first(self.id_field).and_then(|v| v.as_str()).unwrap_or("").to_string(); - let doc_type = retrieved_doc.get_first(self.type_field).and_then(|v| v.as_str()).unwrap_or("").to_string(); - let doc_ns = retrieved_doc.get_first(self.namespace_field).and_then(|v| v.as_str()).unwrap_or(""); - if let Some(ns) = namespace { - if doc_ns != ns && doc_ns != "global" { + let id = retrieved_doc + .get_first(self.id_field) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let doc_type = retrieved_doc + .get_first(self.type_field) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let doc_ns = retrieved_doc + .get_first(self.namespace_field) + .and_then(|v| v.as_str()) + .unwrap_or(""); + if let Some(ns) = namespace + && doc_ns != ns && doc_ns != "global" { continue; } - } results.push((id, doc_type)); } Ok(results) diff --git a/server/src/state.rs b/server/src/state.rs index d7246f7..aa7cf2d 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -1,6 +1,6 @@ use crate::models::*; -use crate::store::{Store, read_json_file, write_json_atomic}; use crate::search::MemoryIndex; +use crate::store::{Store, read_json_file, write_json_atomic}; use std::collections::{HashMap, HashSet}; use std::fs; use std::path::PathBuf; @@ -100,12 +100,15 @@ impl MemoryState { let mut session_graph = self.session_graph.write().unwrap(); update_fn(&mut session_graph); let wal_path = self.base_dir.join("wal.jsonl"); - if let Ok(payload) = serde_json::to_string(&*session_graph) { - if let Ok(mut file) = std::fs::OpenOptions::new().create(true).append(true).open(&wal_path) { + if let Ok(payload) = serde_json::to_string(&*session_graph) + && let Ok(mut file) = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&wal_path) + { use std::io::Write; let _ = writeln!(file, "{}", payload); } - } } pub fn apply_sync_write(&self, update_fn: F) { let lock_path = self.base_dir.join("master.lock"); @@ -142,7 +145,7 @@ impl MemoryState { pub fn rebuild_index(&self) { if let Ok(new_idx) = MemoryIndex::new() { let full = self.get_full_graph(); - for (_, e) in &full.entities { + for e in full.entities.values() { let _ = new_idx.index_entity(e); } for t in self.tasks.read() { @@ -160,5 +163,3 @@ impl MemoryState { } } } - - diff --git a/server/src/store.rs b/server/src/store.rs index 0dd114a..733c8b5 100644 --- a/server/src/store.rs +++ b/server/src/store.rs @@ -1,4 +1,4 @@ -use serde::{de::DeserializeOwned, Serialize}; +use serde::{Serialize, de::DeserializeOwned}; use std::fs; use std::path::{Path, PathBuf}; use std::sync::RwLock; @@ -6,15 +6,17 @@ use std::time::SystemTime; pub fn read_json_file(path: &Path) -> T { if let Ok(data) = fs::read(path) - && let Ok(parsed) = serde_json::from_slice(&data) { - return parsed; - } + && let Ok(parsed) = serde_json::from_slice(&data) + { + return parsed; + } let bak_path = path.with_extension("json.bak"); if let Ok(data) = fs::read(&bak_path) - && let Ok(parsed) = serde_json::from_slice(&data) { - let _ = fs::write(path, data); - return parsed; - } + && let Ok(parsed) = serde_json::from_slice(&data) + { + let _ = fs::write(path, data); + return parsed; + } T::default() } diff --git a/server/src/tools.rs b/server/src/tools.rs index 73a3a36..4282373 100644 --- a/server/src/tools.rs +++ b/server/src/tools.rs @@ -78,7 +78,9 @@ pub struct UpdateTaskStatusTool { pub status: String, } #[derive(Debug, Deserialize, Serialize, JsonSchema)] -pub struct ListActiveTasksTool { pub git_branch: Option, } +pub struct ListActiveTasksTool { + pub git_branch: Option, +} #[derive(Debug, Deserialize, Serialize, JsonSchema)] pub struct StoreSnippetTool { pub name: String, diff --git a/stub/src/main.rs b/stub/src/main.rs index ff7a451..389762c 100644 --- a/stub/src/main.rs +++ b/stub/src/main.rs @@ -3,8 +3,8 @@ use futures_util::StreamExt; use std::sync::Arc; use tokio::io::AsyncBufReadExt; use tokio::sync::RwLock; -use tokio_util::io::StreamReader; use tokio::sync::mpsc; +use tokio_util::io::StreamReader; #[derive(Parser)] #[command(name = "mcp-memory-stub")] @@ -27,7 +27,9 @@ fn main() -> Result<(), Box> { let mut handle = stdin.lock(); let mut buffer = String::new(); while let Ok(bytes) = std::io::BufRead::read_line(&mut handle, &mut buffer) { - if bytes == 0 { break; } + if bytes == 0 { + break; + } let _ = msg_tx.blocking_send(buffer.clone()); buffer.clear(); } @@ -37,7 +39,7 @@ fn main() -> Result<(), Box> { let target_url = cli.target; let post_url = Arc::new(RwLock::new(String::new())); let post_url_clone = Arc::clone(&post_url); - + let client = reqwest::Client::builder().build()?; tokio::spawn(async move { @@ -46,14 +48,16 @@ fn main() -> Result<(), Box> { loop { let url = post_url_clone.read().await.clone(); if !url.is_empty() { - let res = client.post(&url) + let res = client + .post(&url) .header("Accept", "application/json, text/event-stream") .header("Content-Type", "application/json") .body(msg.clone()) - .send().await; - + .send() + .await; + if res.is_ok() { - break; + break; } } tokio::time::sleep(tokio::time::Duration::from_millis(500)).await; @@ -73,20 +77,27 @@ fn main() -> Result<(), Box> { let sse_url = format!("{}/sse", target_url); let client = reqwest::Client::builder().build()?; - - match client.get(&sse_url).header("Accept", "text/event-stream").send().await { + + match client + .get(&sse_url) + .header("Accept", "text/event-stream") + .send() + .await + { Ok(resp) => { if resp.status() == reqwest::StatusCode::GONE { eprintln!("[PROXY] Target gone, exiting."); break; } - - let stream = resp.bytes_stream().map(|res| res.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))); + + let stream = resp.bytes_stream().map(|res| { + res.map_err(std::io::Error::other) + }); let mut reader = tokio::io::BufReader::new(StreamReader::new(stream)); let mut line = String::new(); let mut is_message = false; let mut is_endpoint = false; - + loop { tokio::select! { _ = shutdown_rx.recv() => { @@ -103,14 +114,13 @@ fn main() -> Result<(), Box> { } else if trimmed.starts_with("event: endpoint") { is_endpoint = true; is_message = false; - } else if trimmed.starts_with("data: ") { + } else if let Some(stripped) = trimmed.strip_prefix("data: ") { if is_message { - println!("{}", &trimmed[6..]); + println!("{}", stripped); is_message = false; } else if is_endpoint { - let ep = &trimmed[6..]; let mut p = post_url.write().await; - *p = format!("{}{}", target_url, ep); + *p = format!("{}{}", target_url, stripped); is_endpoint = false; } }