riz/mcp-memory
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: rustfmt, clippy lints and code tidying

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-17 15:35:41 +01:00
1 parent 3716c3e698
commit 1752753fcc
19 files changed
+890 -355
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

linux-nvim/build.rs
+8 -3
View File
@@ -2,19 +2,24 @@ use std::process::Command;
fn main() { fn main() {
let git_hash = Command::new("git") let git_hash = Command::new("git")
.args(&["rev-parse", "--short", "HEAD"]) .args(["rev-parse", "--short", "HEAD"])
.output() .output()
.ok() .ok()
.and_then(|out| String::from_utf8(out.stdout).ok()) .and_then(|out| String::from_utf8(out.stdout).ok())
.unwrap_or_else(|| "unknown".to_string()); .unwrap_or_else(|| "unknown".to_string());
let git_date = Command::new("git") let git_date = Command::new("git")
.args(&["log", "-1", "--format=%cd", "--date=format:%Y.%m.%d"]) .args(["log", "-1", "--format=%cd", "--date=format:%Y.%m.%d"])
.output() .output()
.ok() .ok()
.and_then(|out| String::from_utf8(out.stdout).ok()) .and_then(|out| String::from_utf8(out.stdout).ok())
.unwrap_or_else(|| "unknown".to_string()); .unwrap_or_else(|| "unknown".to_string());
let version = format!("{} ({} {})", env!("CARGO_PKG_VERSION"), git_date.trim(), git_hash.trim()); let version = format!(
"{} ({} {})",
env!("CARGO_PKG_VERSION"),
git_date.trim(),
git_hash.trim()
);
println!("cargo:rustc-env=APP_VERSION={}", version); println!("cargo:rustc-env=APP_VERSION={}", version);
} }
linux-nvim/tests/integration_test.rs
+8 -4
View File
@@ -1,6 +1,5 @@
use serde_json::{json, Value}; use serde_json::Value;
use std::io::{BufRead, BufReader, Read, Write}; use std::io::{BufRead, BufReader, Read, Write};
use std::process::{Command, Stdio};
fn send_message(stdin: &mut std::process::ChildStdin, msg: Value) { fn send_message(stdin: &mut std::process::ChildStdin, msg: Value) {
let s = serde_json::to_string(&msg).unwrap(); let s = serde_json::to_string(&msg).unwrap();
@@ -46,7 +45,10 @@ fn test_mcp_initialization_and_tools_list() {
let mut nvim_exe = std::env::current_exe().unwrap(); let mut nvim_exe = std::env::current_exe().unwrap();
nvim_exe.pop(); nvim_exe.pop();
nvim_exe.pop(); nvim_exe.pop();
nvim_exe.push(format!("mcp-memory-linux-nvim{}", std::env::consts::EXE_SUFFIX)); nvim_exe.push(format!(
"mcp-memory-linux-nvim{}",
std::env::consts::EXE_SUFFIX
));
let mut child = Command::new(&nvim_exe) let mut child = Command::new(&nvim_exe)
.stdin(Stdio::piped()) .stdin(Stdio::piped())
@@ -113,7 +115,9 @@ fn test_mcp_initialization_and_tools_list() {
assert_eq!(tools_resp["jsonrpc"], "2.0"); assert_eq!(tools_resp["jsonrpc"], "2.0");
assert_eq!(tools_resp["id"], 2); assert_eq!(tools_resp["id"], 2);
let tools = tools_resp["result"]["tools"].as_array().expect("result.tools must be an array"); let tools = tools_resp["result"]["tools"]
.as_array()
.expect("result.tools must be an array");
assert!(!tools.is_empty(), "Server must expose at least one tool"); assert!(!tools.is_empty(), "Server must expose at least one tool");
let has_get_active_buffer = tools.iter().any(|t| t["name"] == "nvim_get_active_buffer"); let has_get_active_buffer = tools.iter().any(|t| t["name"] == "nvim_get_active_buffer");
nvim-core/src/lib.rs
+119 -58
View File
@@ -20,7 +20,9 @@ pub struct JsonRpcResponse {
pub error: Option<Value>, pub error: Option<Value>,
} }
pub async fn read_message<R: tokio::io::AsyncRead + Unpin>(stdin: &mut BufReader<R>) -> Option<JsonRpcRequest> { pub async fn read_message<R: tokio::io::AsyncRead + Unpin>(
stdin: &mut BufReader<R>,
) -> Option<JsonRpcRequest> {
let mut length = 0; let mut length = 0;
loop { loop {
let mut line = String::new(); let mut line = String::new();
@@ -32,7 +34,11 @@ pub async fn read_message<R: tokio::io::AsyncRead + Unpin>(stdin: &mut BufReader
return match serde_json::from_str::<JsonRpcRequest>(line.trim_end()) { return match serde_json::from_str::<JsonRpcRequest>(line.trim_end()) {
Ok(req) => Some(req), Ok(req) => Some(req),
Err(e) => { Err(e) => {
tracing::error!("Failed to parse JSON-RPC request from JSONL: {}. Payload: {}", e, line); tracing::error!(
"Failed to parse JSON-RPC request from JSONL: {}. Payload: {}",
e,
line
);
None None
} }
}; };
@@ -58,7 +64,15 @@ pub async fn read_message<R: tokio::io::AsyncRead + Unpin>(stdin: &mut BufReader
pub async fn send_response(response: JsonRpcResponse) { pub async fn send_response(response: JsonRpcResponse) {
let msg = serde_json::to_string(&response).unwrap(); let msg = serde_json::to_string(&response).unwrap();
tracing::info!("Sending JSON-RPC response (id: {:?}): {}", response.id, if msg.len() > 500 { format!("{}...", &msg[..500]) } else { msg.clone() }); tracing::info!(
"Sending JSON-RPC response (id: {:?}): {}",
response.id,
if msg.len() > 500 {
format!("{}...", &msg[..500])
} else {
msg.clone()
}
);
// CRITICAL ARCHITECTURAL DECISION: // CRITICAL ARCHITECTURAL DECISION:
// The MCP StdioTransport MUST use Newline-Delimited JSON (NDJSON). // The MCP StdioTransport MUST use Newline-Delimited JSON (NDJSON).
// Do NOT use LSP-style HTTP headers (e.g. Content-Length). // Do NOT use LSP-style HTTP headers (e.g. Content-Length).
@@ -75,12 +89,14 @@ pub async fn send_error(id: Value, code: i32, message: &str) {
id, id,
result: None, result: None,
error: Some(serde_json::json!({"code": code, "message": message})), error: Some(serde_json::json!({"code": code, "message": message})),
}).await; })
.await;
} }
#[cfg(windows)] #[cfg(windows)]
async fn get_socket_path() -> Result<String, String> { async fn get_socket_path() -> Result<String, String> {
let profile = std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into()); let profile =
std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into());
let path = format!("{}\\.gemini\\active_nvim.txt", profile); let path = format!("{}\\.gemini\\active_nvim.txt", profile);
if let Ok(content) = std::fs::read_to_string(&path) { if let Ok(content) = std::fs::read_to_string(&path) {
@@ -140,12 +156,20 @@ async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
use tokio::net::windows::named_pipe::ClientOptions; use tokio::net::windows::named_pipe::ClientOptions;
let msgid = if let rmpv::Value::Array(ref arr) = req { let msgid = if let rmpv::Value::Array(ref arr) = req {
if arr.len() > 1 { arr[1].clone() } else { rmpv::Value::Nil } if arr.len() > 1 {
} else { rmpv::Value::Nil }; arr[1].clone()
} else {
rmpv::Value::Nil
}
} else {
rmpv::Value::Nil
};
tracing::info!("Connecting to neovim pipe"); tracing::info!("Connecting to neovim pipe");
let socket_path = get_socket_path().await?; let socket_path = get_socket_path().await?;
let mut client = ClientOptions::new().open(&socket_path).map_err(|e| e.to_string())?; let mut client = ClientOptions::new()
.open(&socket_path)
.map_err(|e| e.to_string())?;
let mut buf = Vec::new(); let mut buf = Vec::new();
rmpv::encode::write_value(&mut buf, &req).map_err(|e| e.to_string())?; rmpv::encode::write_value(&mut buf, &req).map_err(|e| e.to_string())?;
@@ -163,18 +187,21 @@ async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
offset += cursor.position() as usize; offset += cursor.position() as usize;
if let rmpv::Value::Array(ref arr) = val { if let rmpv::Value::Array(ref arr) = val {
if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) && arr[1] == msgid { if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) && arr[1] == msgid
{
tracing::info!("Received RPC response from neovim (msgid: {})", msgid); tracing::info!("Received RPC response from neovim (msgid: {})", msgid);
return Ok(val); return Ok(val);
} }
} }
continue; continue;
}, }
Err(_) => { Err(_) => {
let read_future = client.read(&mut chunk); let read_future = client.read(&mut chunk);
match tokio::time::timeout(tokio::time::Duration::from_secs(5), read_future).await { match tokio::time::timeout(tokio::time::Duration::from_secs(5), read_future).await {
Ok(Ok(n)) => { Ok(Ok(n)) => {
if n == 0 { return Err("Connection closed".into()); } if n == 0 {
return Err("Connection closed".into());
}
resp_buf.extend_from_slice(&chunk[..n]); resp_buf.extend_from_slice(&chunk[..n]);
} }
Ok(Err(e)) => return Err(e.to_string()), Ok(Err(e)) => return Err(e.to_string()),
@@ -193,12 +220,20 @@ async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
use tokio::net::UnixStream; use tokio::net::UnixStream;
let msgid = if let rmpv::Value::Array(ref arr) = req { let msgid = if let rmpv::Value::Array(ref arr) = req {
if arr.len() > 1 { arr[1].clone() } else { rmpv::Value::Nil } if arr.len() > 1 {
} else { rmpv::Value::Nil }; arr[1].clone()
} else {
rmpv::Value::Nil
}
} else {
rmpv::Value::Nil
};
tracing::info!("Connecting to neovim socket"); tracing::info!("Connecting to neovim socket");
let socket_path = get_socket_path().await?; let socket_path = get_socket_path().await?;
let mut stream = UnixStream::connect(socket_path).await.map_err(|e| e.to_string())?; let mut stream = UnixStream::connect(socket_path)
.await
.map_err(|e| e.to_string())?;
let mut buf = Vec::new(); let mut buf = Vec::new();
rmpv::encode::write_value(&mut buf, &req).map_err(|e| e.to_string())?; rmpv::encode::write_value(&mut buf, &req).map_err(|e| e.to_string())?;
@@ -216,18 +251,21 @@ async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
offset += cursor.position() as usize; offset += cursor.position() as usize;
if let rmpv::Value::Array(ref arr) = val { if let rmpv::Value::Array(ref arr) = val {
if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) && arr[1] == msgid { if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) && arr[1] == msgid
{
tracing::info!("Received RPC response from neovim (msgid: {})", msgid); tracing::info!("Received RPC response from neovim (msgid: {})", msgid);
return Ok(val); return Ok(val);
} }
} }
continue; continue;
}, }
Err(_) => { Err(_) => {
let read_future = stream.read(&mut chunk); let read_future = stream.read(&mut chunk);
match tokio::time::timeout(tokio::time::Duration::from_secs(5), read_future).await { match tokio::time::timeout(tokio::time::Duration::from_secs(5), read_future).await {
Ok(Ok(n)) => { Ok(Ok(n)) => {
if n == 0 { return Err("Connection closed".into()); } if n == 0 {
return Err("Connection closed".into());
}
resp_buf.extend_from_slice(&chunk[..n]); resp_buf.extend_from_slice(&chunk[..n]);
} }
Ok(Err(e)) => return Err(e.to_string()), Ok(Err(e)) => return Err(e.to_string()),
@@ -300,9 +338,7 @@ async fn get_nvim_cursor() -> Result<String, String> {
RmpValue::Integer(0.into()), RmpValue::Integer(0.into()),
RmpValue::Integer(3.into()), // msgid RmpValue::Integer(3.into()), // msgid
RmpValue::String("nvim_win_get_cursor".into()), RmpValue::String("nvim_win_get_cursor".into()),
RmpValue::Array(vec![ RmpValue::Array(vec![RmpValue::Integer(0.into())]),
RmpValue::Integer(0.into()),
]),
]); ]);
let resp = call_nvim(req).await?; let resp = call_nvim(req).await?;
@@ -360,7 +396,8 @@ async fn get_nvim_visual_selection() -> Result<String, String> {
async fn set_nvim_diagnostics(line: i64, message: &str) -> Result<(), String> { async fn set_nvim_diagnostics(line: i64, message: &str) -> Result<(), String> {
let escaped_message = message.replace("\\", "\\\\").replace("\"", "\\\""); let escaped_message = message.replace("\\", "\\\\").replace("\"", "\\\"");
let lua_script = format!(r#" let lua_script = format!(
r#"
local ns = vim.api.nvim_create_namespace("gemini_diagnostics") local ns = vim.api.nvim_create_namespace("gemini_diagnostics")
local diagnostics = {{{{ local diagnostics = {{{{
lnum = {} - 1, lnum = {} - 1,
@@ -369,7 +406,9 @@ async fn set_nvim_diagnostics(line: i64, message: &str) -> Result<(), String> {
message = "{}", message = "{}",
}}}} }}}}
vim.diagnostic.set(ns, 0, diagnostics, {{}}) vim.diagnostic.set(ns, 0, diagnostics, {{}})
"#, line, escaped_message); "#,
line, escaped_message
);
use rmpv::Value as RmpValue; use rmpv::Value as RmpValue;
let req = RmpValue::Array(vec![ let req = RmpValue::Array(vec![
@@ -404,7 +443,7 @@ fn rmpv_to_json(val: &rmpv::Value) -> serde_json::Value {
} else { } else {
serde_json::Value::Null serde_json::Value::Null
} }
}, }
rmpv::Value::F32(f) => serde_json::json!(f), rmpv::Value::F32(f) => serde_json::json!(f),
rmpv::Value::F64(f) => serde_json::json!(f), rmpv::Value::F64(f) => serde_json::json!(f),
rmpv::Value::String(s) => { rmpv::Value::String(s) => {
@@ -413,11 +452,11 @@ fn rmpv_to_json(val: &rmpv::Value) -> serde_json::Value {
} else { } else {
serde_json::Value::Null serde_json::Value::Null
} }
}, }
rmpv::Value::Array(arr) => { rmpv::Value::Array(arr) => {
let vec: Vec<serde_json::Value> = arr.iter().map(rmpv_to_json).collect(); let vec: Vec<serde_json::Value> = arr.iter().map(rmpv_to_json).collect();
serde_json::Value::Array(vec) serde_json::Value::Array(vec)
}, }
rmpv::Value::Map(map) => { rmpv::Value::Map(map) => {
let mut obj = serde_json::Map::new(); let mut obj = serde_json::Map::new();
for (k, v) in map { for (k, v) in map {
@@ -429,7 +468,7 @@ fn rmpv_to_json(val: &rmpv::Value) -> serde_json::Value {
obj.insert(key_str, rmpv_to_json(v)); obj.insert(key_str, rmpv_to_json(v));
} }
serde_json::Value::Object(obj) serde_json::Value::Object(obj)
}, }
_ => serde_json::json!(format!("{:?}", val)), _ => serde_json::json!(format!("{:?}", val)),
} }
} }
@@ -440,10 +479,7 @@ async fn execute_nvim_lua(code: &str) -> Result<String, String> {
RmpValue::Integer(0.into()), RmpValue::Integer(0.into()),
RmpValue::Integer(6.into()), // msgid RmpValue::Integer(6.into()), // msgid
RmpValue::String("nvim_exec_lua".into()), RmpValue::String("nvim_exec_lua".into()),
RmpValue::Array(vec![ RmpValue::Array(vec![RmpValue::String(code.into()), RmpValue::Array(vec![])]),
RmpValue::String(code.into()),
RmpValue::Array(vec![]),
]),
]); ]);
let resp = call_nvim(req).await?; let resp = call_nvim(req).await?;
@@ -460,7 +496,12 @@ async fn execute_nvim_lua(code: &str) -> Result<String, String> {
} }
pub async fn run_mcp_loop(app_name: &str, app_version: &str) { pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
if std::env::args().any(|arg| arg == "--version") { if std::env::args().any(|arg| arg == "--version") {
println!("{} {} ({})", app_name, app_version, std::env::var("GIT_HASH").unwrap_or_else(|_| "unknown".to_string())); println!(
"{} {} ({})",
app_name,
app_version,
std::env::var("GIT_HASH").unwrap_or_else(|_| "unknown".to_string())
);
return; return;
} }
let _guard = init_logging(app_name); let _guard = init_logging(app_name);
@@ -471,7 +512,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
Some(m) => { Some(m) => {
tracing::info!("Received message method: {}", m.method); tracing::info!("Received message method: {}", m.method);
m m
}, }
None => { None => {
tracing::info!("Stdin closed, exiting loop"); tracing::info!("Stdin closed, exiting loop");
break; break;
@@ -488,9 +529,14 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
match msg.method.as_str() { match msg.method.as_str() {
"initialize" => { "initialize" => {
let init = rmcp::model::InitializeResult::new( let init = rmcp::model::InitializeResult::new(
rmcp::model::ServerCapabilities::builder().enable_tools().build() rmcp::model::ServerCapabilities::builder()
.enable_tools()
.build(),
) )
.with_server_info(rmcp::model::Implementation::new(app_name.clone(), app_version.clone())) .with_server_info(rmcp::model::Implementation::new(
app_name.clone(),
app_version.clone(),
))
.with_protocol_version(rmcp::model::ProtocolVersion::V_2024_11_05); .with_protocol_version(rmcp::model::ProtocolVersion::V_2024_11_05);
send_response(JsonRpcResponse { send_response(JsonRpcResponse {
@@ -498,7 +544,8 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
id, id,
result: Some(serde_json::to_value(init).unwrap()), result: Some(serde_json::to_value(init).unwrap()),
error: None, error: None,
}).await; })
.await;
} }
"notifications/initialized" => {} "notifications/initialized" => {}
"tools/list" => { "tools/list" => {
@@ -595,7 +642,10 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
match name { match name {
"nvim_goto_line" => { "nvim_goto_line" => {
if let (Some(file), Some(line)) = (args.get("file").and_then(|v| v.as_str()), args.get("line").and_then(|v| v.as_i64())) { if let (Some(file), Some(line)) = (
args.get("file").and_then(|v| v.as_str()),
args.get("line").and_then(|v| v.as_i64()),
) {
let escaped_file = file.replace("\\", "\\\\"); let escaped_file = file.replace("\\", "\\\\");
let cmd = format!("e {} | {} | normal! zz", escaped_file, line); let cmd = format!("e {} | {} | normal! zz", escaped_file, line);
match send_nvim_command(&cmd).await { match send_nvim_command(&cmd).await {
@@ -615,8 +665,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
send_error(id, -32602, "Missing 'file' or 'line'").await; send_error(id, -32602, "Missing 'file' or 'line'").await;
} }
} }
"nvim_get_active_buffer" => { "nvim_get_active_buffer" => match get_nvim_active_buffer().await {
match get_nvim_active_buffer().await {
Ok(content) => { Ok(content) => {
send_response(JsonRpcResponse { send_response(JsonRpcResponse {
jsonrpc: "2.0".to_string(), jsonrpc: "2.0".to_string(),
@@ -625,13 +674,12 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
"content": [{"type": "text", "text": content}] "content": [{"type": "text", "text": content}]
})), })),
error: None, error: None,
}).await; })
.await;
} }
Err(e) => send_error(id, -32603, &e).await, Err(e) => send_error(id, -32603, &e).await,
} },
} "nvim_get_cursor" => match get_nvim_cursor().await {
"nvim_get_cursor" => {
match get_nvim_cursor().await {
Ok(content) => { Ok(content) => {
send_response(JsonRpcResponse { send_response(JsonRpcResponse {
jsonrpc: "2.0".to_string(), jsonrpc: "2.0".to_string(),
@@ -640,13 +688,12 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
"content": [{"type": "text", "text": content}] "content": [{"type": "text", "text": content}]
})), })),
error: None, error: None,
}).await; })
.await;
} }
Err(e) => send_error(id, -32603, &e).await, Err(e) => send_error(id, -32603, &e).await,
} },
} "nvim_get_visual_selection" => match get_nvim_visual_selection().await {
"nvim_get_visual_selection" => {
match get_nvim_visual_selection().await {
Ok(content) => { Ok(content) => {
send_response(JsonRpcResponse { send_response(JsonRpcResponse {
jsonrpc: "2.0".to_string(), jsonrpc: "2.0".to_string(),
@@ -655,13 +702,16 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
"content": [{"type": "text", "text": content}] "content": [{"type": "text", "text": content}]
})), })),
error: None, error: None,
}).await; })
.await;
} }
Err(e) => send_error(id, -32603, &e).await, Err(e) => send_error(id, -32603, &e).await,
} },
}
"nvim_set_diagnostics" => { "nvim_set_diagnostics" => {
if let (Some(line), Some(message)) = (args.get("line").and_then(|v| v.as_i64()), args.get("message").and_then(|v| v.as_str())) { if let (Some(line), Some(message)) = (
args.get("line").and_then(|v| v.as_i64()),
args.get("message").and_then(|v| v.as_str()),
) {
match set_nvim_diagnostics(line, message).await { match set_nvim_diagnostics(line, message).await {
Ok(_) => { Ok(_) => {
send_response(JsonRpcResponse { send_response(JsonRpcResponse {
@@ -700,7 +750,8 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
"content": [{"type": "text", "text": result}] "content": [{"type": "text", "text": result}]
})), })),
error: None, error: None,
}).await; })
.await;
} }
Err(e) => send_error(id, -32603, &e).await, Err(e) => send_error(id, -32603, &e).await,
} }
@@ -730,7 +781,8 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
"content": [{"type": "text", "text": result}] "content": [{"type": "text", "text": result}]
})), })),
error: None, error: None,
}).await; })
.await;
} }
Err(e) => send_error(id, -32603, &e).await, Err(e) => send_error(id, -32603, &e).await,
} }
@@ -746,7 +798,8 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
"content": [{"type": "text", "text": result}] "content": [{"type": "text", "text": result}]
})), })),
error: None, error: None,
}).await; })
.await;
} }
Err(e) => send_error(id, -32603, &e).await, Err(e) => send_error(id, -32603, &e).await,
} }
@@ -770,7 +823,9 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
} }
fn init_logging(app_name: &str) -> tracing_appender::non_blocking::WorkerGuard { fn init_logging(app_name: &str) -> tracing_appender::non_blocking::WorkerGuard {
let log_dir = dirs::home_dir().unwrap_or_default().join(".gemini/mcp_memory/logs"); let log_dir = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/logs");
std::fs::create_dir_all(&log_dir).unwrap_or_default(); std::fs::create_dir_all(&log_dir).unwrap_or_default();
let file_appender = tracing_appender::rolling::daily(log_dir, format!("{}.log", app_name)); let file_appender = tracing_appender::rolling::daily(log_dir, format!("{}.log", app_name));
@@ -794,7 +849,10 @@ mod tests {
assert_eq!(rmpv_to_json(&rmpv::Value::Nil), serde_json::Value::Null); assert_eq!(rmpv_to_json(&rmpv::Value::Nil), serde_json::Value::Null);
assert_eq!(rmpv_to_json(&rmpv::Value::Boolean(true)), json!(true)); assert_eq!(rmpv_to_json(&rmpv::Value::Boolean(true)), json!(true));
assert_eq!(rmpv_to_json(&rmpv::Value::Integer(42.into())), json!(42)); assert_eq!(rmpv_to_json(&rmpv::Value::Integer(42.into())), json!(42));
assert_eq!(rmpv_to_json(&rmpv::Value::String("hello".into())), json!("hello")); assert_eq!(
rmpv_to_json(&rmpv::Value::String("hello".into())),
json!("hello")
);
} }
#[test] #[test]
@@ -809,7 +867,10 @@ mod tests {
#[test] #[test]
fn test_rmpv_to_json_map() { fn test_rmpv_to_json_map() {
let mut map = vec![]; let mut map = vec![];
map.push((rmpv::Value::String("key1".into()), rmpv::Value::Integer(100.into()))); map.push((
rmpv::Value::String("key1".into()),
rmpv::Value::Integer(100.into()),
));
let rmp_map = rmpv::Value::Map(map); let rmp_map = rmpv::Value::Map(map);
let json_map = rmpv_to_json(&rmp_map); let json_map = rmpv_to_json(&rmp_map);
server/build.rs
+8 -3
View File
@@ -2,19 +2,24 @@ use std::process::Command;
fn main() { fn main() {
let git_hash = Command::new("git") let git_hash = Command::new("git")
.args(&["rev-parse", "--short", "HEAD"]) .args(["rev-parse", "--short", "HEAD"])
.output() .output()
.ok() .ok()
.and_then(|out| String::from_utf8(out.stdout).ok()) .and_then(|out| String::from_utf8(out.stdout).ok())
.unwrap_or_else(|| "unknown".to_string()); .unwrap_or_else(|| "unknown".to_string());
let git_date = Command::new("git") let git_date = Command::new("git")
.args(&["log", "-1", "--format=%cd", "--date=format:%Y.%m.%d"]) .args(["log", "-1", "--format=%cd", "--date=format:%Y.%m.%d"])
.output() .output()
.ok() .ok()
.and_then(|out| String::from_utf8(out.stdout).ok()) .and_then(|out| String::from_utf8(out.stdout).ok())
.unwrap_or_else(|| "unknown".to_string()); .unwrap_or_else(|| "unknown".to_string());
let version = format!("{} ({} {})", env!("CARGO_PKG_VERSION"), git_date.trim(), git_hash.trim()); let version = format!(
"{} ({} {})",
env!("CARGO_PKG_VERSION"),
git_date.trim(),
git_hash.trim()
);
println!("cargo:rustc-env=APP_VERSION={}", version); println!("cargo:rustc-env=APP_VERSION={}", version);
} }
server/src/bin_test.rs
+5 -3
View File
@@ -1,8 +1,10 @@
use rmcp::model::{InitializeResult, ServerCapabilities}; use rmcp::model::{InitializeResult, ServerCapabilities};
fn main() { fn main() {
let init = InitializeResult::new( let init = InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
ServerCapabilities::builder().enable_tools().build() .with_server_info(rmcp::model::Implementation::new(
).with_server_info(rmcp::model::Implementation::new("gemini-mcp-memory", "3.0.0")); "gemini-mcp-memory",
"3.0.0",
));
println!("{}", serde_json::to_string_pretty(&init).unwrap()); println!("{}", serde_json::to_string_pretty(&init).unwrap());
} }
server/src/handlers.rs
+466 -151
View File
@@ -18,8 +18,6 @@ macro_rules! parse_tool {
}; };
} }
use serde::de::DeserializeOwned; use serde::de::DeserializeOwned;
use std::collections::HashSet; use std::collections::HashSet;
use std::sync::Arc; use std::sync::Arc;
@@ -57,92 +55,276 @@ impl MemoryHandler {
} }
} }
}); });
tracing::debug!("<<< [Server] Replying to server/discover with payload: {}", payload.to_string()); tracing::debug!(
"<<< [Server] Replying to server/discover with payload: {}",
payload.to_string()
);
Some(crate::mcp::success(id, payload)) Some(crate::mcp::success(id, payload))
} }
"initialize" => { "initialize" => {
let init = rmcp::model::InitializeResult::new( let init = rmcp::model::InitializeResult::new(
rmcp::model::ServerCapabilities::builder().enable_tools().build() rmcp::model::ServerCapabilities::builder()
).with_server_info(rmcp::model::Implementation::new("gemini-mcp-memory", "3.0.0")); .enable_tools()
.build(),
)
.with_server_info(rmcp::model::Implementation::new(
"gemini-mcp-memory",
"3.0.0",
));
tracing::debug!("<<< [Server] Replying to initialize with rmcp payload"); tracing::debug!("<<< [Server] Replying to initialize with rmcp payload");
Some(crate::mcp::success(id, serde_json::to_value(&init).unwrap())) Some(crate::mcp::success(
id,
serde_json::to_value(&init).unwrap(),
))
} }
"notifications/initialized" => { "notifications/initialized" => None,
None
}
"tools/list" => { "tools/list" => {
let tools = vec![ let tools = vec![
crate::mcp::tool_def::<crate::tools::QueryGraphPathTool>("query_graph_path", "Traverse the knowledge graph to find a path between two entities."), crate::mcp::tool_def::<crate::tools::QueryGraphPathTool>(
crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entities in the knowledge graph."), "query_graph_path",
crate::mcp::tool_def::<CreateRelationsTool>("create_relations", "Create new relations between entities in the knowledge graph."), "Traverse the knowledge graph to find a path between two entities.",
crate::mcp::tool_def::<AddObservationsTool>("add_observations", "Add new observations to existing entities in the knowledge graph."), ),
crate::mcp::tool_def::<DeleteEntitiesTool>("delete_entities", "Delete entities from the knowledge graph."), crate::mcp::tool_def::<CreateEntitiesTool>(
crate::mcp::tool_def::<DeleteObservationsTool>("delete_observations", "Delete observations from existing entities."), "create_entities",
crate::mcp::tool_def::<DeleteRelationsTool>("delete_relations", "Delete relations between entities."), "Create new entities in the knowledge graph.",
crate::mcp::tool_def::<ReadGraphTool>("read_graph", "Read the entire knowledge graph."), ),
crate::mcp::tool_def::<SearchNodesTool>("search_nodes", "Search for entities in the knowledge graph by name or type."), crate::mcp::tool_def::<CreateRelationsTool>(
crate::mcp::tool_def::<OpenNodesTool>("open_nodes", "Open and retrieve full details of specific nodes in the knowledge graph."), "create_relations",
crate::mcp::tool_def::<LogCodeChangeTool>("log_code_change", "Log a significant code change or refactor in the memory system."), "Create new relations between entities in the knowledge graph.",
crate::mcp::tool_def::<QueryRecentChangesTool>("query_recent_changes", "Query recently logged code changes."), ),
crate::mcp::tool_def::<VisualizeGraphTool>("visualize_graph", "Generate a visual representation of the knowledge graph."), crate::mcp::tool_def::<AddObservationsTool>(
crate::mcp::tool_def::<AddStickyNoteTool>("add_sticky_note", "Add a sticky note for unstructured thoughts or reminders."), "add_observations",
crate::mcp::tool_def::<ReadStickyNotesTool>("read_sticky_notes", "Read all active sticky notes."), "Add new observations to existing entities in the knowledge graph.",
crate::mcp::tool_def::<CondenseEntityTool>("condense_entity", "Condense or summarize an entity's observations to reduce size."), ),
crate::mcp::tool_def::<AddTaskTool>("add_task", "Add a new task to the task tracker."), crate::mcp::tool_def::<DeleteEntitiesTool>(
crate::mcp::tool_def::<UpdateTaskStatusTool>("update_task_status", "Update the status of an existing task."), "delete_entities",
crate::mcp::tool_def::<DeleteTaskTool>("delete_task", "Delete a task and all its children."), "Delete entities from the knowledge graph.",
crate::mcp::tool_def::<ListActiveTasksTool>("list_active_tasks", "List all currently active tasks."), ),
crate::mcp::tool_def::<SetAcceptanceCriteriaTool>("set_acceptance_criteria", "Define a strict checklist of acceptance criteria for a given task."), crate::mcp::tool_def::<DeleteObservationsTool>(
crate::mcp::tool_def::<VerifyAcceptanceCriteriaTool>("verify_acceptance_criteria", "Mark a previously defined acceptance criteria as met."), "delete_observations",
crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Store a reusable code snippet."), "Delete observations from existing entities.",
crate::mcp::tool_def::<SearchSnippetsTool>("search_snippets", "Search through stored code snippets."), ),
crate::mcp::tool_def::<DeleteSnippetTool>("delete_snippet", "Delete a stored code snippet."), crate::mcp::tool_def::<DeleteRelationsTool>(
crate::mcp::tool_def::<LogDecisionTool>("log_decision", "Log an architectural decision record (ADR)."), "delete_relations",
crate::mcp::tool_def::<QueryDecisionsTool>("query_decisions", "Query architectural decision records."), "Delete relations between entities.",
crate::mcp::tool_def::<MergeEntitiesTool>("merge_entities", "Merge two entities in the knowledge graph into one."), ),
crate::mcp::tool_def::<FindOrphansTool>("find_orphans", "Find orphaned entities (entities without any relations) in the graph."), crate::mcp::tool_def::<ReadGraphTool>(
crate::mcp::tool_def::<LearnPreferenceTool>("learn_preference", "Record a user preference or behavior to adapt future interactions."), "read_graph",
crate::mcp::tool_def::<ReadPreferencesTool>("read_preferences", "Read all learned user preferences."), "Read the entire knowledge graph.",
crate::mcp::tool_def::<LogErrorFixTool>("log_error_fix", "Log a complex error and its fix for future reference."), ),
crate::mcp::tool_def::<SearchErrorFixesTool>("search_error_fixes", "Search through previously logged error fixes."), crate::mcp::tool_def::<SearchNodesTool>(
crate::mcp::tool_def::<PinFileTool>("pin_file", "Pin a file to keep it explicitly in the context workspace."), "search_nodes",
crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Unpin a file from the context workspace."), "Search for entities in the knowledge graph by name or type.",
crate::mcp::tool_def::<ListPinnedFilesTool>("list_pinned_files", "List all currently pinned files."), ),
crate::mcp::tool_def::<AddSessionSummaryTool>("add_session_summary", "Add a summary of the current session."), crate::mcp::tool_def::<OpenNodesTool>(
crate::mcp::tool_def::<GetProjectTimelineTool>("get_project_timeline", "Get a timeline of major project events."), "open_nodes",
crate::mcp::tool_def::<LeaveHandoffMemoTool>("leave_handoff_memo", "Leave a memo for the next session or agent."), "Open and retrieve full details of specific nodes in the knowledge graph.",
crate::mcp::tool_def::<ReadHandoffMemosTool>("read_handoff_memos", "Read pending handoff memos."), ),
crate::mcp::tool_def::<ClearHandoffMemosTool>("clear_handoff_memos", "Clear handoff memos after reading."), crate::mcp::tool_def::<LogCodeChangeTool>(
crate::mcp::tool_def::<UpdateEnvFingerprintTool>("update_env_fingerprint", "Update the environment fingerprint (e.g., OS, tool versions)."), "log_code_change",
crate::mcp::tool_def::<ReadEnvFingerprintTool>("read_env_fingerprint", "Read the current environment fingerprint."), "Log a significant code change or refactor in the memory system.",
crate::mcp::tool_def::<LogEnvRequirementTool>("log_env_requirement", "Log a required tool or package for the environment."), ),
crate::mcp::tool_def::<AddMilestoneTool>("add_milestone", "Add a new project milestone."), crate::mcp::tool_def::<QueryRecentChangesTool>(
crate::mcp::tool_def::<UpdateMilestoneTool>("update_milestone", "Update the status of a project milestone."), "query_recent_changes",
crate::mcp::tool_def::<ListMilestonesTool>("list_milestones", "List all project milestones."), "Query recently logged code changes.",
),
crate::mcp::tool_def::<VisualizeGraphTool>(
"visualize_graph",
"Generate a visual representation of the knowledge graph.",
),
crate::mcp::tool_def::<AddStickyNoteTool>(
"add_sticky_note",
"Add a sticky note for unstructured thoughts or reminders.",
),
crate::mcp::tool_def::<ReadStickyNotesTool>(
"read_sticky_notes",
"Read all active sticky notes.",
),
crate::mcp::tool_def::<CondenseEntityTool>(
"condense_entity",
"Condense or summarize an entity's observations to reduce size.",
),
crate::mcp::tool_def::<AddTaskTool>(
"add_task",
"Add a new task to the task tracker.",
),
crate::mcp::tool_def::<UpdateTaskStatusTool>(
"update_task_status",
"Update the status of an existing task.",
),
crate::mcp::tool_def::<DeleteTaskTool>(
"delete_task",
"Delete a task and all its children.",
),
crate::mcp::tool_def::<ListActiveTasksTool>(
"list_active_tasks",
"List all currently active tasks.",
),
crate::mcp::tool_def::<SetAcceptanceCriteriaTool>(
"set_acceptance_criteria",
"Define a strict checklist of acceptance criteria for a given task.",
),
crate::mcp::tool_def::<VerifyAcceptanceCriteriaTool>(
"verify_acceptance_criteria",
"Mark a previously defined acceptance criteria as met.",
),
crate::mcp::tool_def::<StoreSnippetTool>(
"store_snippet",
"Store a reusable code snippet.",
),
crate::mcp::tool_def::<SearchSnippetsTool>(
"search_snippets",
"Search through stored code snippets.",
),
crate::mcp::tool_def::<DeleteSnippetTool>(
"delete_snippet",
"Delete a stored code snippet.",
),
crate::mcp::tool_def::<LogDecisionTool>(
"log_decision",
"Log an architectural decision record (ADR).",
),
crate::mcp::tool_def::<QueryDecisionsTool>(
"query_decisions",
"Query architectural decision records.",
),
crate::mcp::tool_def::<MergeEntitiesTool>(
"merge_entities",
"Merge two entities in the knowledge graph into one.",
),
crate::mcp::tool_def::<FindOrphansTool>(
"find_orphans",
"Find orphaned entities (entities without any relations) in the graph.",
),
crate::mcp::tool_def::<LearnPreferenceTool>(
"learn_preference",
"Record a user preference or behavior to adapt future interactions.",
),
crate::mcp::tool_def::<ReadPreferencesTool>(
"read_preferences",
"Read all learned user preferences.",
),
crate::mcp::tool_def::<LogErrorFixTool>(
"log_error_fix",
"Log a complex error and its fix for future reference.",
),
crate::mcp::tool_def::<SearchErrorFixesTool>(
"search_error_fixes",
"Search through previously logged error fixes.",
),
crate::mcp::tool_def::<PinFileTool>(
"pin_file",
"Pin a file to keep it explicitly in the context workspace.",
),
crate::mcp::tool_def::<UnpinFileTool>(
"unpin_file",
"Unpin a file from the context workspace.",
),
crate::mcp::tool_def::<ListPinnedFilesTool>(
"list_pinned_files",
"List all currently pinned files.",
),
crate::mcp::tool_def::<AddSessionSummaryTool>(
"add_session_summary",
"Add a summary of the current session.",
),
crate::mcp::tool_def::<GetProjectTimelineTool>(
"get_project_timeline",
"Get a timeline of major project events.",
),
crate::mcp::tool_def::<LeaveHandoffMemoTool>(
"leave_handoff_memo",
"Leave a memo for the next session or agent.",
),
crate::mcp::tool_def::<ReadHandoffMemosTool>(
"read_handoff_memos",
"Read pending handoff memos.",
),
crate::mcp::tool_def::<ClearHandoffMemosTool>(
"clear_handoff_memos",
"Clear handoff memos after reading.",
),
crate::mcp::tool_def::<UpdateEnvFingerprintTool>(
"update_env_fingerprint",
"Update the environment fingerprint (e.g., OS, tool versions).",
),
crate::mcp::tool_def::<ReadEnvFingerprintTool>(
"read_env_fingerprint",
"Read the current environment fingerprint.",
),
crate::mcp::tool_def::<LogEnvRequirementTool>(
"log_env_requirement",
"Log a required tool or package for the environment.",
),
crate::mcp::tool_def::<AddMilestoneTool>(
"add_milestone",
"Add a new project milestone.",
),
crate::mcp::tool_def::<UpdateMilestoneTool>(
"update_milestone",
"Update the status of a project milestone.",
),
crate::mcp::tool_def::<ListMilestonesTool>(
"list_milestones",
"List all project milestones.",
),
crate::mcp::tool_def::<GenerateStandupReportTool>( crate::mcp::tool_def::<GenerateStandupReportTool>(
"generate_standup_report", "generate_standup_report",
"", "",
), ),
crate::mcp::tool_def::<RegisterEnvironmentTool>("register_environment", "Register details about a specific deployment environment."), crate::mcp::tool_def::<RegisterEnvironmentTool>(
"register_environment",
"Register details about a specific deployment environment.",
),
crate::mcp::tool_def::<GetEnvironmentDetailsTool>( crate::mcp::tool_def::<GetEnvironmentDetailsTool>(
"get_environment_details", "get_environment_details",
"", "",
), ),
crate::mcp::tool_def::<AddPrChecklistItemTool>("add_pr_checklist_item", "Add an item to the PR checklist."), crate::mcp::tool_def::<AddPrChecklistItemTool>(
crate::mcp::tool_def::<GetPrChecklistTool>("get_pr_checklist", "Get the current PR checklist."), "add_pr_checklist_item",
crate::mcp::tool_def::<ClearPrChecklistTool>("clear_pr_checklist", "Clear the PR checklist."), "Add an item to the PR checklist.",
crate::mcp::tool_def::<LogTechDebtTool>("log_tech_debt", "Log identified technical debt."), ),
crate::mcp::tool_def::<ResolveTechDebtTool>("resolve_tech_debt", "Mark a logged technical debt as resolved."), crate::mcp::tool_def::<GetPrChecklistTool>(
crate::mcp::tool_def::<ListTechDebtTool>("list_tech_debt", "List all unresolved technical debt."), "get_pr_checklist",
crate::mcp::tool_def::<SaveContextWorkspaceTool>("save_context_workspace", "Save the current set of pinned files and context."), "Get the current PR checklist.",
crate::mcp::tool_def::<LoadContextWorkspaceTool>("load_context_workspace", "Load a previously saved context workspace."), ),
crate::mcp::tool_def::<ClearPrChecklistTool>(
"clear_pr_checklist",
"Clear the PR checklist.",
),
crate::mcp::tool_def::<LogTechDebtTool>(
"log_tech_debt",
"Log identified technical debt.",
),
crate::mcp::tool_def::<ResolveTechDebtTool>(
"resolve_tech_debt",
"Mark a logged technical debt as resolved.",
),
crate::mcp::tool_def::<ListTechDebtTool>(
"list_tech_debt",
"List all unresolved technical debt.",
),
crate::mcp::tool_def::<SaveContextWorkspaceTool>(
"save_context_workspace",
"Save the current set of pinned files and context.",
),
crate::mcp::tool_def::<LoadContextWorkspaceTool>(
"load_context_workspace",
"Load a previously saved context workspace.",
),
crate::mcp::tool_def::<ListContextWorkspacesTool>( crate::mcp::tool_def::<ListContextWorkspacesTool>(
"list_context_workspaces", "list_context_workspaces",
"", "",
), ),
crate::mcp::tool_def::<OmniSearchTool>("omni_search", "Search across all memory sources (graph, tasks, snippets, ADRs, etc.) at once."), crate::mcp::tool_def::<OmniSearchTool>(
crate::mcp::tool_def::<GetProjectHealthTool>("get_project_health", "Get a synthesized health report of the project based on memory data."), "omni_search",
"Search across all memory sources (graph, tasks, snippets, ADRs, etc.) at once.",
),
crate::mcp::tool_def::<GetProjectHealthTool>(
"get_project_health",
"Get a synthesized health report of the project based on memory data.",
),
]; ];
Some(crate::mcp::success( Some(crate::mcp::success(
id, id,
@@ -157,7 +339,8 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
.cloned() .cloned()
.unwrap_or(serde_json::Value::Object(Default::default())); .unwrap_or(serde_json::Value::Object(Default::default()));
self.state.broadcast_activity(&format!("Agent executed tool: {}", name)); self.state
.broadcast_activity(&format!("Agent executed tool: {}", name));
let result: Result<String, String> = match name { let result: Result<String, String> = match name {
"query_graph_path" => { "query_graph_path" => {
@@ -166,7 +349,8 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
let max_depth = req.max_depth.unwrap_or(5); let max_depth = req.max_depth.unwrap_or(5);
let mut queue = std::collections::VecDeque::new(); let mut queue = std::collections::VecDeque::new();
let mut visited = std::collections::HashSet::new(); let mut visited = std::collections::HashSet::new();
let mut parents: std::collections::HashMap<String, (String, String)> = std::collections::HashMap::new(); let mut parents: std::collections::HashMap<String, (String, String)> =
std::collections::HashMap::new();
queue.push_back(req.start_node.clone()); queue.push_back(req.start_node.clone());
visited.insert(req.start_node.clone()); visited.insert(req.start_node.clone());
@@ -186,12 +370,21 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
for rel in &graph.relations { for rel in &graph.relations {
if rel.from == current && !visited.contains(&rel.to) { if rel.from == current && !visited.contains(&rel.to) {
visited.insert(rel.to.clone()); visited.insert(rel.to.clone());
parents.insert(rel.to.clone(), (current.clone(), rel.relation_type.clone())); parents.insert(
rel.to.clone(),
(current.clone(), rel.relation_type.clone()),
);
queue.push_back(rel.to.clone()); queue.push_back(rel.to.clone());
nodes_at_next_depth += 1; nodes_at_next_depth += 1;
} else if rel.to == current && !visited.contains(&rel.from) { } else if rel.to == current && !visited.contains(&rel.from) {
visited.insert(rel.from.clone()); visited.insert(rel.from.clone());
parents.insert(rel.from.clone(), (current.clone(), format!("inverse({})", rel.relation_type))); parents.insert(
rel.from.clone(),
(
current.clone(),
format!("inverse({})", rel.relation_type),
),
);
queue.push_back(rel.from.clone()); queue.push_back(rel.from.clone());
nodes_at_next_depth += 1; nodes_at_next_depth += 1;
} }
@@ -215,7 +408,10 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
path.reverse(); path.reverse();
Ok(format!("Path found:\n{}", path.join("\n"))) Ok(format!("Path found:\n{}", path.join("\n")))
} else { } else {
Ok(format!("No path found between {} and {} within depth {}", req.start_node, req.end_node, max_depth)) Ok(format!(
"No path found between {} and {} within depth {}",
req.start_node, req.end_node, max_depth
))
} }
} }
"create_entities" => { "create_entities" => {
@@ -248,8 +444,7 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
let full = self.state.get_full_graph(); let full = self.state.get_full_graph();
self.state.write_to_local_delta(|g| { self.state.write_to_local_delta(|g| {
for o in req.observations { for o in req.observations {
if let Some(full_e) = full.entities.get(&o.entity_name) if let Some(full_e) = full.entities.get(&o.entity_name) {
{
let mut e = let mut e =
g.entities.get(&o.entity_name).cloned().unwrap_or_else( g.entities.get(&o.entity_name).cloned().unwrap_or_else(
|| Entity { || Entity {
@@ -284,8 +479,7 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
let req = parse_tool!(args.clone(), id, DeleteObservationsTool); let req = parse_tool!(args.clone(), id, DeleteObservationsTool);
self.state.apply_sync_write(|master| { self.state.apply_sync_write(|master| {
for d in req.deletions { for d in req.deletions {
if let Some(e) = master.entities.get_mut(&d.entity_name) if let Some(e) = master.entities.get_mut(&d.entity_name) {
{
let to_rem: HashSet<_> = d.observations.into_iter().collect(); let to_rem: HashSet<_> = d.observations.into_iter().collect();
e.observations.retain(|o| !to_rem.contains(o)); e.observations.retain(|o| !to_rem.contains(o));
} }
@@ -335,7 +529,8 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
let full = self.state.get_full_graph(); let full = self.state.get_full_graph();
for (id, doc_type) in matches { for (id, doc_type) in matches {
if doc_type == "entity" if doc_type == "entity"
&& let Some(e) = full.entities.get(&id) { && let Some(e) = full.entities.get(&id)
{
result.entities.insert(id, e.clone()); result.entities.insert(id, e.clone());
} }
} }
@@ -502,7 +697,7 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
created_at: now, created_at: now,
updated_at: now, updated_at: now,
git_branch: req.git_branch, git_branch: req.git_branch,
parent_id: parent_id, parent_id,
dependencies: deps, dependencies: deps,
acceptance_criteria: vec![], acceptance_criteria: vec![],
}; };
@@ -527,21 +722,24 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
while added_new { while added_new {
added_new = false; added_new = false;
for t in tasks.iter() { for t in tasks.iter() {
if let Some(pid) = &t.parent_id { if let Some(pid) = &t.parent_id
if to_delete.contains(pid) && !to_delete.contains(&t.id) { && to_delete.contains(pid) && !to_delete.contains(&t.id) {
to_delete.insert(t.id.clone()); to_delete.insert(t.id.clone());
added_new = true; added_new = true;
} }
} }
} }
}
tasks.retain(|t| !to_delete.contains(&t.id)); tasks.retain(|t| !to_delete.contains(&t.id));
deleted_count = initial_len - tasks.len(); deleted_count = initial_len - tasks.len();
}); });
if deleted_count > 0 { if deleted_count > 0 {
Ok(vec![format!("Deleted task and its children ({} total).", deleted_count).to_string()][0].clone()) Ok(vec![
format!("Deleted task and its children ({} total).", deleted_count)
.to_string(),
][0]
.clone())
} else { } else {
Ok(vec!["Task not found.".to_string()][0].clone()) Ok(vec!["Task not found.".to_string()][0].clone())
} }
@@ -556,20 +754,24 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
self.state.tasks.modify(|tasks| { self.state.tasks.modify(|tasks| {
// Find target task // Find target task
let mut target_id = String::new(); let mut target_id = String::new();
if let Some(t) = tasks.iter().find(|t| t.id == req.id || t.title == req.id) { if let Some(t) =
tasks.iter().find(|t| t.id == req.id || t.title == req.id)
{
target_id = t.id.clone(); target_id = t.id.clone();
} }
if target_id.is_empty() { return; } if target_id.is_empty() {
return;
}
found = true; found = true;
if target_status == "done" || target_status == "completed" { if target_status == "done" || target_status == "completed" {
// 1. Check Acceptance Criteria // 1. Check Acceptance Criteria
if let Some(t) = tasks.iter().find(|t| t.id == target_id) { if let Some(t) = tasks.iter().find(|t| t.id == target_id)
if t.acceptance_criteria.iter().any(|c| !c.is_met) { && t.acceptance_criteria.iter().any(|c| !c.is_met) {
blocked = true; blocked = true;
blocker_details = "Unmet acceptance criteria exist.".to_string(); blocker_details =
} "Unmet acceptance criteria exist.".to_string();
} }
// 2. Check dependencies // 2. Check dependencies
@@ -577,30 +779,41 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
let mut uncompleted_deps = Vec::new(); let mut uncompleted_deps = Vec::new();
if let Some(t) = tasks.iter().find(|t| t.id == target_id) { if let Some(t) = tasks.iter().find(|t| t.id == target_id) {
for dep_id in &t.dependencies { for dep_id in &t.dependencies {
if let Some(dep_task) = tasks.iter().find(|dt| dt.id == *dep_id) { if let Some(dep_task) =
if dep_task.status != "completed" && dep_task.status != "done" { tasks.iter().find(|dt| dt.id == *dep_id)
&& dep_task.status != "completed"
&& dep_task.status != "done"
{
uncompleted_deps.push(dep_task.title.clone()); uncompleted_deps.push(dep_task.title.clone());
} }
} }
} }
}
if !uncompleted_deps.is_empty() { if !uncompleted_deps.is_empty() {
blocked = true; blocked = true;
blocker_details = format!("Blocked by dependencies: {}", uncompleted_deps.join(", ")); blocker_details = format!(
"Blocked by dependencies: {}",
uncompleted_deps.join(", ")
);
} }
} }
// 3. Check child tasks // 3. Check child tasks
if !blocked { if !blocked {
let mut uncompleted_children = Vec::new(); let mut uncompleted_children = Vec::new();
for child in tasks.iter().filter(|t| t.parent_id.as_ref() == Some(&target_id)) { for child in tasks
.iter()
.filter(|t| t.parent_id.as_ref() == Some(&target_id))
{
if child.status != "completed" && child.status != "done" { if child.status != "completed" && child.status != "done" {
uncompleted_children.push(child.title.clone()); uncompleted_children.push(child.title.clone());
} }
} }
if !uncompleted_children.is_empty() { if !uncompleted_children.is_empty() {
blocked = true; blocked = true;
blocker_details = format!("Blocked by child tasks: {}", uncompleted_children.join(", ")); blocker_details = format!(
"Blocked by child tasks: {}",
uncompleted_children.join(", ")
);
} }
} }
} }
@@ -609,7 +822,10 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
// Apply update // Apply update
if let Some(t) = tasks.iter_mut().find(|t| t.id == target_id) { if let Some(t) = tasks.iter_mut().find(|t| t.id == target_id) {
t.status = target_status.clone(); t.status = target_status.clone();
t.updated_at = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(); t.updated_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
} }
// Cascade cancellation to children // Cascade cancellation to children
@@ -619,7 +835,9 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
while i < to_cancel.len() { while i < to_cancel.len() {
let current_pid = to_cancel[i].clone(); let current_pid = to_cancel[i].clone();
for t in tasks.iter_mut() { for t in tasks.iter_mut() {
if t.parent_id.as_ref() == Some(&current_pid) && t.status != "completed" { if t.parent_id.as_ref() == Some(&current_pid)
&& t.status != "completed"
{
t.status = target_status.clone(); t.status = target_status.clone();
to_cancel.push(t.id.clone()); to_cancel.push(t.id.clone());
} }
@@ -631,7 +849,13 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
}); });
if blocked { if blocked {
Ok(vec![format!("Error: Cannot transition task. {}", blocker_details)].into_iter().next().unwrap()) Ok(vec![format!(
"Error: Cannot transition task. {}",
blocker_details
)]
.into_iter()
.next()
.unwrap())
} else if found { } else if found {
Ok(vec!["Task status updated.".to_string()][0].clone()) Ok(vec!["Task status updated.".to_string()][0].clone())
} else { } else {
@@ -655,18 +879,30 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
let req = parse_tool!(args.clone(), id, SetAcceptanceCriteriaTool); let req = parse_tool!(args.clone(), id, SetAcceptanceCriteriaTool);
let mut success = false; let mut success = false;
self.state.tasks.modify(|tasks| { self.state.tasks.modify(|tasks| {
if let Some(task) = tasks.iter_mut().rev().find(|t| t.title == req.task_title) { if let Some(task) =
task.acceptance_criteria = req.criteria.into_iter().map(|desc| crate::models::AcceptanceCriteria { tasks.iter_mut().rev().find(|t| t.title == req.task_title)
{
task.acceptance_criteria = req
.criteria
.into_iter()
.map(|desc| crate::models::AcceptanceCriteria {
id: uuid::Uuid::new_v4().to_string(), id: uuid::Uuid::new_v4().to_string(),
description: desc, description: desc,
is_met: false, is_met: false,
}).collect(); })
task.updated_at = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(); .collect();
task.updated_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
success = true; success = true;
} }
}); });
if success { if success {
Ok(vec!["Acceptance criteria set successfully.".to_string()][0].clone()) Ok(
vec!["Acceptance criteria set successfully.".to_string()][0]
.clone(),
)
} else { } else {
Ok(vec!["Task not found.".to_string()][0].clone()) Ok(vec!["Task not found.".to_string()][0].clone())
} }
@@ -676,24 +912,37 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
let mut success = false; let mut success = false;
let mut already_met = false; let mut already_met = false;
self.state.tasks.modify(|tasks| { self.state.tasks.modify(|tasks| {
if let Some(task) = tasks.iter_mut().find(|t| t.id == req.task_id) { if let Some(task) = tasks.iter_mut().find(|t| t.id == req.task_id)
if let Some(ac) = task.acceptance_criteria.iter_mut().find(|c| c.id == req.criteria || c.description == req.criteria) { && let Some(ac) = task
.acceptance_criteria
.iter_mut()
.find(|c| c.id == req.criteria || c.description == req.criteria)
{
if ac.is_met { if ac.is_met {
already_met = true; already_met = true;
} else { } else {
ac.is_met = true; ac.is_met = true;
success = true; success = true;
task.updated_at = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(); task.updated_at = SystemTime::now()
} .duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
} }
} }
}); });
if success { if success {
Ok(vec![format!("Acceptance criteria verified with proof: {}", req.proof)][0].clone()) Ok(vec![format!(
"Acceptance criteria verified with proof: {}",
req.proof
)][0]
.clone())
} else if already_met { } else if already_met {
Ok(vec!["Acceptance criteria was already met.".to_string()][0].clone()) Ok(vec!["Acceptance criteria was already met.".to_string()][0].clone())
} else { } else {
Ok(vec!["Acceptance criteria or task not found.".to_string()][0].clone()) Ok(
vec!["Acceptance criteria or task not found.".to_string()][0]
.clone(),
)
} }
} }
"store_snippet" => { "store_snippet" => {
@@ -1224,7 +1473,8 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
let full = self.state.get_full_graph(); let full = self.state.get_full_graph();
for (id, doc_type) in &matches { for (id, doc_type) in &matches {
if doc_type == "entity" if doc_type == "entity"
&& let Some(e) = full.entities.get(id) { && let Some(e) = full.entities.get(id)
{
kg.entities.insert(id.clone(), e.clone()); kg.entities.insert(id.clone(), e.clone());
} }
} }
@@ -1345,23 +1595,19 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
}; };
match result { match result {
Ok(text) => { Ok(text) => Some(crate::mcp::success(
Some(crate::mcp::success(
id, id,
serde_json::json!({ serde_json::json!({
"content": [{ "type": "text", "text": text }] "content": [{ "type": "text", "text": text }]
}), }),
)) )),
} Err(e) => Some(crate::mcp::success(
Err(e) => {
Some(crate::mcp::success(
id, id,
serde_json::json!({ serde_json::json!({
"isError": true, "isError": true,
"content": [{ "type": "text", "text": e }] "content": [{ "type": "text", "text": e }]
}), }),
)) )),
}
} }
} }
_ => { _ => {
@@ -1381,13 +1627,19 @@ crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Create new entiti
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use std::sync::Arc;
use crate::state::MemoryState; use crate::state::MemoryState;
use serde_json::json; use serde_json::json;
use std::sync::Arc;
#[tokio::test] #[tokio::test]
async fn test_handle_initialize() { async fn test_handle_initialize() {
let store_dir = std::env::temp_dir().join(format!("mcp_test_handlers_{}", std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs())); let store_dir = std::env::temp_dir().join(format!(
"mcp_test_handlers_{}",
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs()
));
std::fs::create_dir_all(&store_dir).unwrap(); std::fs::create_dir_all(&store_dir).unwrap();
let redb_path = store_dir.join("mcp_store.redb"); let redb_path = store_dir.join("mcp_store.redb");
let db = Arc::new(redb::Database::create(&redb_path).unwrap()); let db = Arc::new(redb::Database::create(&redb_path).unwrap());
@@ -1401,7 +1653,9 @@ mod tests {
let state = Arc::new(MemoryState { let state = Arc::new(MemoryState {
base_dir: store_dir.clone(), base_dir: store_dir.clone(),
graph: crate::store::Store::new("knowledge_graph_master", db.clone()), graph: crate::store::Store::new("knowledge_graph_master", db.clone()),
search_index: std::sync::RwLock::new(crate::search::MemoryIndex::new(&store_dir).unwrap()), search_index: std::sync::RwLock::new(
crate::search::MemoryIndex::new(&store_dir).unwrap(),
),
ledger: crate::store::Store::new("audit_ledger", db.clone()), ledger: crate::store::Store::new("audit_ledger", db.clone()),
sticky: crate::store::Store::new("sticky_notes", db.clone()), sticky: crate::store::Store::new("sticky_notes", db.clone()),
tasks: crate::store::Store::new("tasks", db.clone()), tasks: crate::store::Store::new("tasks", db.clone()),
@@ -1419,7 +1673,8 @@ mod tests {
pr_checklists: crate::store::Store::new("pr_checklists", db.clone()), pr_checklists: crate::store::Store::new("pr_checklists", db.clone()),
tech_debts: crate::store::Store::new("tech_debts", db.clone()), tech_debts: crate::store::Store::new("tech_debts", db.clone()),
gates: crate::store::Store::new("gates", db.clone()), gates: crate::store::Store::new("gates", db.clone()),
context_workspaces: crate::store::Store::new("context_workspaces", db.clone()), activity_tx: tokio::sync::broadcast::channel(100).0, context_workspaces: crate::store::Store::new("context_workspaces", db.clone()),
activity_tx: tokio::sync::broadcast::channel(100).0,
}); });
let handler = MemoryHandler { state }; let handler = MemoryHandler { state };
@@ -1437,7 +1692,10 @@ mod tests {
} }
}); });
let response = handler.handle_request(req).await.expect("Expected a response"); let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 1); assert_eq!(response["id"], 1);
assert!(response.get("result").is_some()); assert!(response.get("result").is_some());
@@ -1452,7 +1710,11 @@ mod tests {
} }
fn setup_test_handler(test_name: &str) -> MemoryHandler { fn setup_test_handler(test_name: &str) -> MemoryHandler {
let store_dir = std::env::temp_dir().join(format!("mcp_test_handlers_{}_{}", test_name, uuid::Uuid::new_v4())); let store_dir = std::env::temp_dir().join(format!(
"mcp_test_handlers_{}_{}",
test_name,
uuid::Uuid::new_v4()
));
std::fs::create_dir_all(&store_dir).unwrap(); std::fs::create_dir_all(&store_dir).unwrap();
let redb_path = store_dir.join("mcp_store.redb"); let redb_path = store_dir.join("mcp_store.redb");
let db = Arc::new(redb::Database::create(&redb_path).unwrap()); let db = Arc::new(redb::Database::create(&redb_path).unwrap());
@@ -1466,7 +1728,9 @@ mod tests {
base_dir: store_dir.clone(), base_dir: store_dir.clone(),
graph: crate::store::Store::new("knowledge_graph_master", db.clone()), graph: crate::store::Store::new("knowledge_graph_master", db.clone()),
search_index: std::sync::RwLock::new(crate::search::MemoryIndex::new(&store_dir).unwrap()), search_index: std::sync::RwLock::new(
crate::search::MemoryIndex::new(&store_dir).unwrap(),
),
ledger: crate::store::Store::new("audit_ledger", db.clone()), ledger: crate::store::Store::new("audit_ledger", db.clone()),
sticky: crate::store::Store::new("sticky_notes", db.clone()), sticky: crate::store::Store::new("sticky_notes", db.clone()),
tasks: crate::store::Store::new("tasks", db.clone()), tasks: crate::store::Store::new("tasks", db.clone()),
@@ -1484,7 +1748,8 @@ mod tests {
pr_checklists: crate::store::Store::new("pr_checklists", db.clone()), pr_checklists: crate::store::Store::new("pr_checklists", db.clone()),
tech_debts: crate::store::Store::new("tech_debts", db.clone()), tech_debts: crate::store::Store::new("tech_debts", db.clone()),
gates: crate::store::Store::new("gates", db.clone()), gates: crate::store::Store::new("gates", db.clone()),
context_workspaces: crate::store::Store::new("context_workspaces", db.clone()), activity_tx: tokio::sync::broadcast::channel(100).0, context_workspaces: crate::store::Store::new("context_workspaces", db.clone()),
activity_tx: tokio::sync::broadcast::channel(100).0,
}); });
MemoryHandler { state } MemoryHandler { state }
} }
@@ -1500,15 +1765,26 @@ mod tests {
"params": {} "params": {}
}); });
let response = handler.handle_request(req).await.expect("Expected a response"); let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 2); assert_eq!(response["id"], 2);
let tools = response["result"]["tools"].as_array().expect("Tools must be an array"); let tools = response["result"]["tools"]
.as_array()
.expect("Tools must be an array");
assert!(!tools.is_empty()); assert!(!tools.is_empty());
// Verify a specific tool is registered // Verify a specific tool is registered
let add_task_tool = tools.iter().find(|t| t["name"] == "add_task").expect("add_task tool missing"); let add_task_tool = tools
assert_eq!(add_task_tool["description"], "Add a new task to the task tracker."); .iter()
.find(|t| t["name"] == "add_task")
.expect("add_task tool missing");
assert_eq!(
add_task_tool["description"],
"Add a new task to the task tracker."
);
} }
#[tokio::test] #[tokio::test]
@@ -1529,12 +1805,20 @@ mod tests {
} }
}); });
let response = handler.handle_request(req).await.expect("Expected a response"); let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 3); assert_eq!(response["id"], 3);
let content = &response["result"]["content"][0]; let content = &response["result"]["content"][0];
assert_eq!(content["type"], "text"); assert_eq!(content["type"], "text");
assert!(content["text"].as_str().unwrap().starts_with("Task added with ID: ")); assert!(
content["text"]
.as_str()
.unwrap()
.starts_with("Task added with ID: ")
);
// Verify task was actually added to store // Verify task was actually added to store
let tasks = handler.state.tasks.read(); let tasks = handler.state.tasks.read();
@@ -1566,7 +1850,10 @@ mod tests {
} }
}); });
let response = handler.handle_request(req).await.expect("Expected a response"); let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 4); assert_eq!(response["id"], 4);
let content = &response["result"]["content"][0]; let content = &response["result"]["content"][0];
@@ -1574,7 +1861,10 @@ mod tests {
// Verify entity was actually added to state // Verify entity was actually added to state
let session_graph = handler.state.graph.read(); let session_graph = handler.state.graph.read();
let entity = session_graph.entities.get("MemoryHandler").expect("Entity should be in session graph"); let entity = session_graph
.entities
.get("MemoryHandler")
.expect("Entity should be in session graph");
assert_eq!(entity.entity_type, "struct"); assert_eq!(entity.entity_type, "struct");
assert_eq!(entity.observations, vec!["Handles MCP requests natively"]); assert_eq!(entity.observations, vec!["Handles MCP requests natively"]);
assert_eq!(entity.namespace, "core".to_string()); assert_eq!(entity.namespace, "core".to_string());
@@ -1599,7 +1889,10 @@ mod tests {
} }
}); });
let response = handler.handle_request(req).await.expect("Expected a response"); let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 5); assert_eq!(response["id"], 5);
let snippets = handler.state.snippets.read(); let snippets = handler.state.snippets.read();
@@ -1624,7 +1917,10 @@ mod tests {
} }
}); });
let response = handler.handle_request(req).await.expect("Expected a response"); let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 6); assert_eq!(response["id"], 6);
let notes = handler.state.sticky.read(); let notes = handler.state.sticky.read();
@@ -1666,13 +1962,16 @@ mod tests {
let handler = setup_test_handler("add_observations"); let handler = setup_test_handler("add_observations");
// Pre-populate entity // Pre-populate entity
handler.state.graph.modify(|session| { handler.state.graph.modify(|session| {
session.entities.insert("NodeA".to_string(), crate::models::Entity { session.entities.insert(
"NodeA".to_string(),
crate::models::Entity {
name: "NodeA".to_string(), name: "NodeA".to_string(),
entity_type: "class".to_string(), entity_type: "class".to_string(),
observations: vec!["Initial".to_string()], observations: vec!["Initial".to_string()],
namespace: "".to_string(), namespace: "".to_string(),
git_branch: None, git_branch: None,
}); },
);
}); });
let req = json!({ let req = json!({
"jsonrpc": "2.0", "jsonrpc": "2.0",
@@ -1700,13 +1999,16 @@ mod tests {
async fn test_handle_delete_entities() { async fn test_handle_delete_entities() {
let handler = setup_test_handler("delete_entities"); let handler = setup_test_handler("delete_entities");
handler.state.graph.modify(|session| { handler.state.graph.modify(|session| {
session.entities.insert("ToDelete".to_string(), crate::models::Entity { session.entities.insert(
"ToDelete".to_string(),
crate::models::Entity {
name: "ToDelete".to_string(), name: "ToDelete".to_string(),
entity_type: "var".to_string(), entity_type: "var".to_string(),
observations: vec![], observations: vec![],
namespace: "".to_string(), namespace: "".to_string(),
git_branch: None, git_branch: None,
}); },
);
}); });
// Force flush session to master // Force flush session to master
handler.state.apply_sync_write(|_| {}); handler.state.apply_sync_write(|_| {});
@@ -1731,13 +2033,16 @@ mod tests {
async fn test_handle_delete_observations() { async fn test_handle_delete_observations() {
let handler = setup_test_handler("delete_observations"); let handler = setup_test_handler("delete_observations");
handler.state.graph.modify(|session| { handler.state.graph.modify(|session| {
session.entities.insert("NodeA".to_string(), crate::models::Entity { session.entities.insert(
"NodeA".to_string(),
crate::models::Entity {
name: "NodeA".to_string(), name: "NodeA".to_string(),
entity_type: "class".to_string(), entity_type: "class".to_string(),
observations: vec!["Keep".to_string(), "Drop".to_string()], observations: vec!["Keep".to_string(), "Drop".to_string()],
namespace: "".to_string(), namespace: "".to_string(),
git_branch: None, git_branch: None,
}); },
);
}); });
handler.state.apply_sync_write(|_| {}); handler.state.apply_sync_write(|_| {});
let req = json!({ let req = json!({
@@ -1797,7 +2102,10 @@ mod tests {
description: "".to_string(), description: "".to_string(),
created_at: 0, created_at: 0,
updated_at: 0, updated_at: 0,
git_branch: None, acceptance_criteria: vec![], dependencies: vec![], parent_id: None, git_branch: None,
acceptance_criteria: vec![],
dependencies: vec![],
parent_id: None,
}); });
tasks.push(crate::models::Task { tasks.push(crate::models::Task {
id: "2".to_string(), id: "2".to_string(),
@@ -1806,7 +2114,10 @@ mod tests {
description: "".to_string(), description: "".to_string(),
created_at: 0, created_at: 0,
updated_at: 0, updated_at: 0,
git_branch: None, acceptance_criteria: vec![], dependencies: vec![], parent_id: None, git_branch: None,
acceptance_criteria: vec![],
dependencies: vec![],
parent_id: None,
}); });
}); });
@@ -1927,13 +2238,16 @@ mod tests {
async fn test_handle_read_graph() { async fn test_handle_read_graph() {
let handler = setup_test_handler("read_graph"); let handler = setup_test_handler("read_graph");
handler.state.graph.modify(|session| { handler.state.graph.modify(|session| {
session.entities.insert("NodeA".to_string(), crate::models::Entity { session.entities.insert(
"NodeA".to_string(),
crate::models::Entity {
name: "NodeA".to_string(), name: "NodeA".to_string(),
entity_type: "var".to_string(), entity_type: "var".to_string(),
observations: vec![], observations: vec![],
namespace: "".to_string(), namespace: "".to_string(),
git_branch: None, git_branch: None,
}); },
);
}); });
handler.state.apply_sync_write(|_| {}); handler.state.apply_sync_write(|_| {});
@@ -1962,7 +2276,9 @@ mod tests {
git_branch: None, git_branch: None,
}; };
handler.state.graph.modify(|session| { handler.state.graph.modify(|session| {
session.entities.insert("UserRepository".to_string(), entity); session
.entities
.insert("UserRepository".to_string(), entity);
}); });
handler.state.apply_sync_write(|_| {}); handler.state.apply_sync_write(|_| {});
@@ -2322,7 +2638,10 @@ mod tests {
description: "".to_string(), description: "".to_string(),
created_at: 0, created_at: 0,
updated_at: 0, updated_at: 0,
git_branch: None, acceptance_criteria: vec![], dependencies: vec![], parent_id: None, git_branch: None,
acceptance_criteria: vec![],
dependencies: vec![],
parent_id: None,
}); });
}); });
@@ -2343,7 +2662,3 @@ mod tests {
assert_eq!(tasks[0].status, "in_progress"); assert_eq!(tasks[0].status, "in_progress");
} }
} }
server/src/main.rs
+129 -65
View File
@@ -1,4 +1,7 @@
#![cfg_attr(not(target_os = "windows"), allow(dead_code, unused_imports, unreachable_code))] #![cfg_attr(
not(target_os = "windows"),
allow(dead_code, unused_imports, unreachable_code)
)]
mod handlers; mod handlers;
mod mcp; mod mcp;
@@ -13,11 +16,11 @@ use crate::models::*;
use crate::state::MemoryState; use crate::state::MemoryState;
use crate::store::Store; use crate::store::Store;
use redb::ReadableTable;
use std::fs; use std::fs;
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::{Arc, RwLock}; use std::sync::{Arc, RwLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH}; use std::time::{Duration, SystemTime, UNIX_EPOCH};
use redb::ReadableTable;
use tokio::time::sleep; use tokio::time::sleep;
use clap::{Parser, Subcommand}; use clap::{Parser, Subcommand};
@@ -83,7 +86,6 @@ enum GateCommands {
}, },
} }
async fn reconcile_worker(state: Arc<MemoryState>) { async fn reconcile_worker(state: Arc<MemoryState>) {
loop { loop {
sleep(Duration::from_secs(5)).await; sleep(Duration::from_secs(5)).await;
@@ -106,14 +108,18 @@ async fn reconcile_worker(state: Arc<MemoryState>) {
let state_clone = state.clone(); let state_clone = state.clone();
let _ = tokio::task::spawn_blocking(move || { let _ = tokio::task::spawn_blocking(move || {
state_clone.rebuild_index(); state_clone.rebuild_index();
}).await; })
.await;
} }
} }
} }
use axum::{ use axum::{
Json, Router, Json, Router,
extract::{Query, State, ws::{WebSocket, Message}}, extract::{
Query, State,
ws::{Message, WebSocket},
},
response::IntoResponse, response::IntoResponse,
routing::{get, post}, routing::{get, post},
}; };
@@ -186,9 +192,11 @@ async fn gate_verify_handler(
(axum::http::StatusCode::FORBIDDEN, msg).into_response() (axum::http::StatusCode::FORBIDDEN, msg).into_response()
} }
} }
None => { None => (
(axum::http::StatusCode::NOT_FOUND, "Action not yet authorized (no gate record found).").into_response() axum::http::StatusCode::NOT_FOUND,
} "Action not yet authorized (no gate record found).",
)
.into_response(),
} }
} }
@@ -229,18 +237,23 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
rt.block_on(async { rt.block_on(async {
tokio::spawn(reconcile_worker(Arc::clone(&state))); tokio::spawn(reconcile_worker(Arc::clone(&state)));
let app_state = Arc::new(AppState { let app_state = Arc::new(AppState {
handler: Arc::new(MemoryHandler { state: Arc::clone(&state) }), handler: Arc::new(MemoryHandler {
state: Arc::clone(&state),
}),
clients: RwLock::new(HashMap::new()), clients: RwLock::new(HashMap::new()),
next_id: AtomicUsize::new(1), next_id: AtomicUsize::new(1),
}); });
let app = Router::new() let app = Router::new()
.route("/api/version", get(|| async move { .route(
"/api/version",
get(|| async move {
axum::Json(serde_json::json!({ axum::Json(serde_json::json!({
"version": env!("APP_VERSION"), "version": env!("APP_VERSION"),
"git_hash": option_env!("GIT_HASH").unwrap_or("unknown") "git_hash": option_env!("GIT_HASH").unwrap_or("unknown")
})) }))
})) }),
)
.route("/ws", get(ws_handler)) .route("/ws", get(ws_handler))
.route("/health", get(health_handler)) .route("/health", get(health_handler))
.route("/nvim/telemetry", post(nvim_telemetry_handler)) .route("/nvim/telemetry", post(nvim_telemetry_handler))
@@ -260,15 +273,19 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
"/", "/",
get(|| async move { axum::response::Html(include_str!("dashboard.html")) }), get(|| async move { axum::response::Html(include_str!("dashboard.html")) }),
) )
.route("/api/graph", get({ .route(
"/api/graph",
get({
let state_clone = app_state.handler.state.clone(); let state_clone = app_state.handler.state.clone();
move || async move { move || async move {
let graph = state_clone.get_full_graph(); let graph = state_clone.get_full_graph();
axum::Json(graph) axum::Json(graph)
} }
})) }),
)
.route("/api/tasks/{id}/complete", post({ .route(
"/api/tasks/{id}/complete",
post({
let state_clone = app_state.handler.state.clone(); let state_clone = app_state.handler.state.clone();
move |axum::extract::Path(id): axum::extract::Path<String>| async move { move |axum::extract::Path(id): axum::extract::Path<String>| async move {
state_clone.tasks.modify(|tasks| { state_clone.tasks.modify(|tasks| {
@@ -281,28 +298,38 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
}); });
axum::Json(serde_json::json!({"status": "success"})) axum::Json(serde_json::json!({"status": "success"}))
} }
})) }),
)
.route("/api/tasks", get({ .route(
"/api/tasks",
get({
let state_clone = app_state.handler.state.clone(); let state_clone = app_state.handler.state.clone();
move || async move { move || async move {
let tasks = state_clone.tasks.read(); let tasks = state_clone.tasks.read();
axum::Json(tasks.clone()) axum::Json(tasks.clone())
} }
})) }),
.route("/api/sticky", get({ )
.route(
"/api/sticky",
get({
let state_clone = app_state.handler.state.clone(); let state_clone = app_state.handler.state.clone();
move || async move { move || async move {
let sticky = state_clone.sticky.read(); let sticky = state_clone.sticky.read();
axum::Json(sticky.clone()) axum::Json(sticky.clone())
} }
})) }),
.route("/api/search", get({ )
.route(
"/api/search",
get({
let state_clone = app_state.handler.state.clone(); let state_clone = app_state.handler.state.clone();
move |axum::extract::Query(params): axum::extract::Query<std::collections::HashMap<String, String>>| async move { move |axum::extract::Query(params): axum::extract::Query<
if let Some(q) = params.get("q") { std::collections::HashMap<String, String>,
if let Ok(idx) = state_clone.search_index.read() { >| async move {
if let Ok(results) = idx.search(q, None) { if let Some(q) = params.get("q")
&& let Ok(idx) = state_clone.search_index.read()
&& let Ok(results) = idx.search(q, None) {
let mut formatted_results = Vec::new(); let mut formatted_results = Vec::new();
for (type_name, content) in results { for (type_name, content) in results {
formatted_results.push(serde_json::json!({ formatted_results.push(serde_json::json!({
@@ -311,14 +338,16 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
"score": 1.0 "score": 1.0
})); }));
} }
return axum::Json(serde_json::json!({ "results": formatted_results })); return axum::Json(
} serde_json::json!({ "results": formatted_results }),
} );
} }
axum::Json(serde_json::json!({ "results": [] })) axum::Json(serde_json::json!({ "results": [] }))
} }
})) }),
.route("/api/stats", )
.route(
"/api/stats",
get({ get({
let state_clone = app_state.handler.state.clone(); let state_clone = app_state.handler.state.clone();
move || async move { move || async move {
@@ -386,29 +415,39 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
}) })
} }
async fn ws_handler( async fn ws_handler(
ws: axum::extract::ws::WebSocketUpgrade, ws: axum::extract::ws::WebSocketUpgrade,
headers: axum::http::HeaderMap, _headers: axum::http::HeaderMap,
axum::extract::State(state): axum::extract::State<Arc<AppState>>, axum::extract::State(state): axum::extract::State<Arc<AppState>>,
axum::extract::Query(query): axum::extract::Query<std::collections::HashMap<String, String>>, axum::extract::Query(query): axum::extract::Query<std::collections::HashMap<String, String>>,
) -> axum::response::Response { ) -> axum::response::Response {
let client_type = query.get("client").cloned().unwrap_or_else(|| "unknown".to_string()); let client_type = query
ws.on_upgrade(move |socket| handle_socket(socket, state, client_type)).into_response() .get("client")
.cloned()
.unwrap_or_else(|| "unknown".to_string());
ws.on_upgrade(move |socket| handle_socket(socket, state, client_type))
.into_response()
} }
async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: String) { async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: String) {
let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst)); let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst));
let (tx, mut rx) = mpsc::channel::<String>(100); let (tx, mut rx) = mpsc::channel::<String>(100);
state.clients.write().unwrap().insert(session_id.clone(), tx.clone()); state
.clients
.write()
.unwrap()
.insert(session_id.clone(), tx.clone());
let (mut sender, mut receiver) = socket.split(); let (mut sender, mut receiver) = socket.split();
let mut send_task = tokio::spawn(async move { let mut send_task = tokio::spawn(async move {
while let Some(msg) = rx.recv().await { while let Some(msg) = rx.recv().await {
tracing::trace!("Sending message to websocket (length: {}): {}", msg.len(), msg); tracing::trace!(
"Sending message to websocket (length: {}): {}",
msg.len(),
msg
);
if sender.send(Message::Text(msg.into())).await.is_err() { if sender.send(Message::Text(msg.into())).await.is_err() {
tracing::error!("Failed to send message to websocket"); tracing::error!("Failed to send message to websocket");
break; break;
@@ -426,14 +465,21 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
while let Some(msg_result) = receiver.next().await { while let Some(msg_result) = receiver.next().await {
match msg_result { match msg_result {
Ok(Message::Text(text)) => { Ok(Message::Text(text)) => {
tracing::info!("Received text message from websocket (length: {})", text.len()); tracing::info!(
"Received text message from websocket (length: {})",
text.len()
);
tracing::trace!("Message content: {}", text); tracing::trace!("Message content: {}", text);
if let Ok(payload) = serde_json::from_str::<serde_json::Value>(&text) { if let Ok(payload) = serde_json::from_str::<serde_json::Value>(&text) {
if client_type == "proxy" { if client_type == "proxy" {
// Send activity broadcast to UI clients // Send activity broadcast to UI clients
if let Some(method) = payload.get("method").and_then(|m| m.as_str()) { if let Some(method) = payload.get("method").and_then(|m| m.as_str())
if method == "tools/call" { && method == "tools/call" {
let name = payload.get("params").and_then(|p| p.get("name")).and_then(|n| n.as_str()).unwrap_or("unknown_tool"); let name = payload
.get("params")
.and_then(|p| p.get("name"))
.and_then(|n| n.as_str())
.unwrap_or("unknown_tool");
let activity_msg = format!("Agent executed tool: {}", name); let activity_msg = format!("Agent executed tool: {}", name);
let event = serde_json::json!({ let event = serde_json::json!({
@@ -448,24 +494,39 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
} }
} }
} }
}
} // End if proxy } // End if proxy
// Process MCP request // Process MCP request
if let Some(response) = handler.handle_request(payload).await { if let Some(response) = handler.handle_request(payload).await {
let res_str = serde_json::to_string(&response).unwrap(); let res_str = serde_json::to_string(&response).unwrap();
let tx_opt = state_clone.clients.read().unwrap().get(&session_id_clone).cloned(); let tx_opt = state_clone
.clients
.read()
.unwrap()
.get(&session_id_clone)
.cloned();
if let Some(client_tx) = tx_opt { if let Some(client_tx) = tx_opt {
if let Err(e) = client_tx.send(res_str).await { if let Err(e) = client_tx.send(res_str).await {
tracing::error!("Failed to send response to client channel for session {}: {}", session_id_clone, e); tracing::error!(
"Failed to send response to client channel for session {}: {}",
session_id_clone,
e
);
} }
} else { } else {
tracing::warn!("Could not find client_tx for session_id {} when trying to send response", session_id_clone); tracing::warn!(
"Could not find client_tx for session_id {} when trying to send response",
session_id_clone
);
} }
} }
} // End if let Ok(payload) }
// End if let Ok(payload)
else { else {
tracing::warn!("Failed to parse payload as JSON from websocket message: {}", text); tracing::warn!(
"Failed to parse payload as JSON from websocket message: {}",
text
);
} }
} // End Ok(Message::Text(text)) } // End Ok(Message::Text(text))
Ok(other) => { Ok(other) => {
@@ -477,7 +538,10 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
} }
} }
} }
tracing::info!("Websocket receiver task ended for session {}", session_id_clone); tracing::info!(
"Websocket receiver task ended for session {}",
session_id_clone
);
}); });
tokio::select! { tokio::select! {
@@ -492,10 +556,12 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
}; };
state.clients.write().unwrap().remove(&session_id); state.clients.write().unwrap().remove(&session_id);
tracing::info!("Websocket session {} closed and removed from state", session_id); tracing::info!(
"Websocket session {} closed and removed from state",
session_id
);
} }
#[derive(serde::Deserialize, serde::Serialize, Debug)] #[derive(serde::Deserialize, serde::Serialize, Debug)]
pub struct NvimTelemetry { pub struct NvimTelemetry {
pub session_id: String, pub session_id: String,
@@ -510,8 +576,10 @@ async fn nvim_telemetry_handler(
axum::Json(payload): axum::Json<NvimTelemetry>, axum::Json(payload): axum::Json<NvimTelemetry>,
) -> impl axum::response::IntoResponse { ) -> impl axum::response::IntoResponse {
// 1. Update active_nvim.txt if FocusGained, BufEnter, or VimEnter // 1. Update active_nvim.txt if FocusGained, BufEnter, or VimEnter
if payload.event == "FocusGained" || payload.event == "BufEnter" || payload.event == "VimEnter" { if payload.event == "FocusGained" || payload.event == "BufEnter" || payload.event == "VimEnter"
let profile = std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into()); {
let profile =
std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into());
let win_path = format!("{}\\.gemini\\active_nvim.txt", profile); let win_path = format!("{}\\.gemini\\active_nvim.txt", profile);
let _ = std::fs::write(&win_path, &payload.session_id); let _ = std::fs::write(&win_path, &payload.session_id);
@@ -527,7 +595,7 @@ async fn nvim_telemetry_handler(
let msg_str = ws_msg.to_string(); let msg_str = ws_msg.to_string();
let clients = state.clients.read().unwrap().clone(); let clients = state.clients.read().unwrap().clone();
for (_, tx) in clients.iter() { for tx in clients.values() {
let _ = tx.send(msg_str.clone()).await; let _ = tx.send(msg_str.clone()).await;
} }
@@ -611,7 +679,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
} }
} }
let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| { let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
dirs::home_dir() dirs::home_dir()
.map(|mut h| { .map(|mut h| {
@@ -657,13 +724,14 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
for (key, file_name) in stores.iter() { for (key, file_name) in stores.iter() {
if table.get(*key).unwrap().is_none() { if table.get(*key).unwrap().is_none() {
let json_path = base.join(file_name); let json_path = base.join(file_name);
if json_path.exists() { if json_path.exists()
if let Ok(data) = fs::read(&json_path) { && let Ok(data) = fs::read(&json_path)
if serde_json::from_slice::<serde_json::Value>(&data).is_ok() { && serde_json::from_slice::<serde_json::Value>(&data).is_ok() {
table.insert(*key, data.as_slice()).unwrap(); table.insert(*key, data.as_slice()).unwrap();
let _ = fs::rename(&json_path, json_path.with_extension("json.migrated")); let _ = fs::rename(
} &json_path,
} json_path.with_extension("json.migrated"),
);
} }
} }
} }
@@ -700,7 +768,3 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
run_server(state) run_server(state)
} }
server/src/models.rs
+1 -1
View File
@@ -1,6 +1,6 @@
use schemars::JsonSchema;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::collections::HashMap; use std::collections::HashMap;
use schemars::JsonSchema;
#[derive(Debug, Clone, Serialize, Deserialize, Default)] #[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct CodeChange { pub struct CodeChange {
server/src/search.rs
+9 -4
View File
@@ -28,7 +28,8 @@ impl MemoryIndex {
let index_dir = store_dir.join("tantivy_index"); let index_dir = store_dir.join("tantivy_index");
std::fs::create_dir_all(&index_dir).unwrap(); 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 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 writer = index.writer(50_000_000)?;
let reader = index let reader = index
@@ -72,7 +73,6 @@ impl MemoryIndex {
Ok(()) Ok(())
} }
pub fn commit(&self) -> tantivy::Result<()> { pub fn commit(&self) -> tantivy::Result<()> {
let mut writer = self.writer.lock().unwrap(); let mut writer = self.writer.lock().unwrap();
writer.commit()?; writer.commit()?;
@@ -113,7 +113,9 @@ impl MemoryIndex {
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
.unwrap_or(""); .unwrap_or("");
if let Some(ns) = namespace if let Some(ns) = namespace
&& doc_ns != ns && doc_ns != "global" { && doc_ns != ns
&& doc_ns != "global"
{
continue; continue;
} }
results.push((id, doc_type)); results.push((id, doc_type));
@@ -172,7 +174,10 @@ mod tests {
status: "open".to_string(), status: "open".to_string(),
created_at: 0, created_at: 0,
updated_at: 0, updated_at: 0,
git_branch: None, acceptance_criteria: vec![], dependencies: vec![], parent_id: None, git_branch: None,
acceptance_criteria: vec![],
dependencies: vec![],
parent_id: None,
}; };
index.index_task(&task).unwrap(); index.index_task(&task).unwrap();
server/src/state.rs
+7 -5
View File
@@ -33,19 +33,21 @@ pub struct MemoryState {
impl MemoryState { impl MemoryState {
pub fn unique_items<T: Eq + std::hash::Hash + Clone>(input: Vec<T>) -> Vec<T> { pub fn unique_items<T: Eq + std::hash::Hash + Clone>(input: Vec<T>) -> Vec<T> {
let mut keys = std::collections::HashSet::new(); let mut keys = std::collections::HashSet::new();
input.into_iter().filter(|entry| keys.insert(entry.clone())).collect() input
.into_iter()
.filter(|entry| keys.insert(entry.clone()))
.collect()
} }
pub fn broadcast_activity(&self, message: &str) { pub fn broadcast_activity(&self, message: &str) {
let payload = serde_json::json!({ let payload = serde_json::json!({
"type": "activity", "type": "activity",
"data": message "data": message
}).to_string(); })
.to_string();
let _ = self.activity_tx.send(payload); let _ = self.activity_tx.send(payload);
} }
pub fn get_full_graph(&self) -> KnowledgeGraph { pub fn get_full_graph(&self) -> KnowledgeGraph {
self.graph.read() self.graph.read()
} }
@@ -61,7 +63,7 @@ impl MemoryState {
pub fn rebuild_index(&self) { pub fn rebuild_index(&self) {
if let Ok(new_idx) = MemoryIndex::new(&self.base_dir) { if let Ok(new_idx) = MemoryIndex::new(&self.base_dir) {
let graph = self.graph.read(); let graph = self.graph.read();
for (_, e) in &graph.entities { for e in graph.entities.values() {
let _ = new_idx.index_entity(e); let _ = new_idx.index_entity(e);
} }
server/src/store.rs
+17 -7
View File
@@ -22,13 +22,11 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + 'static> Store<T
fn load_from_db(key: &str, db: &Database) -> T { fn load_from_db(key: &str, db: &Database) -> T {
let read_txn = db.begin_read().unwrap(); let read_txn = db.begin_read().unwrap();
if let Ok(table) = read_txn.open_table(STORE_TABLE) { if let Ok(table) = read_txn.open_table(STORE_TABLE)
if let Ok(Some(value)) = table.get(key) { && let Ok(Some(value)) = table.get(key)
if let Ok(parsed) = serde_json::from_slice::<T>(value.value()) { && let Ok(parsed) = serde_json::from_slice::<T>(value.value()) {
return parsed; return parsed;
} }
}
}
T::default() T::default()
} }
@@ -94,11 +92,23 @@ mod tests {
// Need to wait for spawn_blocking to finish // Need to wait for spawn_blocking to finish
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
assert_eq!(store.read(), TestData { name: "Hello".to_string(), value: 42 }); assert_eq!(
store.read(),
TestData {
name: "Hello".to_string(),
value: 42
}
);
// Load again to verify persistence // Load again to verify persistence
let store2 = Store::<TestData>::new("test_key", db.clone()); let store2 = Store::<TestData>::new("test_key", db.clone());
assert_eq!(store2.read(), TestData { name: "Hello".to_string(), value: 42 }); assert_eq!(
store2.read(),
TestData {
name: "Hello".to_string(),
value: 42
}
);
} }
#[tokio::test(flavor = "multi_thread", worker_threads = 4)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
server/tests/parity_test.rs
+28 -13
View File
@@ -3,37 +3,44 @@ use std::collections::HashSet;
#[test] #[test]
fn test_eager_tools_parity() { fn test_eager_tools_parity() {
// 1. Read handlers.rs to get memory tools // 1. Read handlers.rs to get memory tools
let memory_source = std::fs::read_to_string("src/handlers.rs").expect("Failed to read handlers.rs"); let memory_source =
std::fs::read_to_string("src/handlers.rs").expect("Failed to read handlers.rs");
let mut memory_tools = HashSet::new(); let mut memory_tools = HashSet::new();
for line in memory_source.lines() { for line in memory_source.lines() {
if line.contains("crate::mcp::tool_def") { if line.contains("crate::mcp::tool_def")
if let Some(start) = line.find("(\"") { && let Some(start) = line.find("(\"") {
let rest = &line[start + 2..]; let rest = &line[start + 2..];
if let Some(end) = rest.find("\"") { if let Some(end) = rest.find("\"") {
memory_tools.insert(rest[..end].to_string()); memory_tools.insert(rest[..end].to_string());
} }
} }
} }
} assert!(
assert!(!memory_tools.is_empty(), "Could not find memory tools in handlers.rs"); !memory_tools.is_empty(),
"Could not find memory tools in handlers.rs"
);
// 2. Read nvim-core/src/lib.rs to get nvim tools // 2. Read nvim-core/src/lib.rs to get nvim tools
let nvim_source = std::fs::read_to_string("../nvim-core/src/lib.rs").expect("Failed to read nvim lib.rs"); let nvim_source =
std::fs::read_to_string("../nvim-core/src/lib.rs").expect("Failed to read nvim lib.rs");
let mut nvim_tools = HashSet::new(); let mut nvim_tools = HashSet::new();
for line in nvim_source.lines() { for line in nvim_source.lines() {
if line.contains("\"name\": \"nvim_") { if line.contains("\"name\": \"nvim_")
if let Some(start) = line.find("\"name\": \"") { && let Some(start) = line.find("\"name\": \"") {
let rest = &line[start + 9..]; let rest = &line[start + 9..];
if let Some(end) = rest.find("\"") { if let Some(end) = rest.find("\"") {
nvim_tools.insert(rest[..end].to_string()); nvim_tools.insert(rest[..end].to_string());
} }
} }
} }
} assert!(
assert!(!nvim_tools.is_empty(), "Could not find nvim tools in lib.rs"); !nvim_tools.is_empty(),
"Could not find nvim tools in lib.rs"
);
// 3. Read Windows mcp_config.json // 3. Read Windows mcp_config.json
let win_home = std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".to_string()); let win_home =
std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".to_string());
let win_config_path = std::path::PathBuf::from(win_home).join(".gemini/config/mcp_config.json"); let win_config_path = std::path::PathBuf::from(win_home).join(".gemini/config/mcp_config.json");
if win_config_path.exists() { if win_config_path.exists() {
let config_str = std::fs::read_to_string(&win_config_path).unwrap(); let config_str = std::fs::read_to_string(&win_config_path).unwrap();
@@ -42,14 +49,22 @@ fn test_eager_tools_parity() {
if let Some(eager) = config["mcpServers"]["memory"]["eagerTools"].as_array() { if let Some(eager) = config["mcpServers"]["memory"]["eagerTools"].as_array() {
for tool in eager { for tool in eager {
let name = tool.as_str().unwrap(); let name = tool.as_str().unwrap();
assert!(memory_tools.contains(name), "Windows config Memory tool '{}' not implemented in handlers.rs!", name); assert!(
memory_tools.contains(name),
"Windows config Memory tool '{}' not implemented in handlers.rs!",
name
);
} }
} }
if let Some(nvim_eager) = config["mcpServers"]["win-nvim"]["eagerTools"].as_array() { if let Some(nvim_eager) = config["mcpServers"]["win-nvim"]["eagerTools"].as_array() {
for tool in nvim_eager { for tool in nvim_eager {
let name = tool.as_str().unwrap(); let name = tool.as_str().unwrap();
assert!(nvim_tools.contains(name), "Windows config Nvim tool '{}' not implemented in nvim-core!", name); assert!(
nvim_tools.contains(name),
"Windows config Nvim tool '{}' not implemented in nvim-core!",
name
);
} }
} }
} }
stub/build.rs
+8 -3
View File
@@ -2,19 +2,24 @@ use std::process::Command;
fn main() { fn main() {
let git_hash = Command::new("git") let git_hash = Command::new("git")
.args(&["rev-parse", "--short", "HEAD"]) .args(["rev-parse", "--short", "HEAD"])
.output() .output()
.ok() .ok()
.and_then(|out| String::from_utf8(out.stdout).ok()) .and_then(|out| String::from_utf8(out.stdout).ok())
.unwrap_or_else(|| "unknown".to_string()); .unwrap_or_else(|| "unknown".to_string());
let git_date = Command::new("git") let git_date = Command::new("git")
.args(&["log", "-1", "--format=%cd", "--date=format:%Y.%m.%d"]) .args(["log", "-1", "--format=%cd", "--date=format:%Y.%m.%d"])
.output() .output()
.ok() .ok()
.and_then(|out| String::from_utf8(out.stdout).ok()) .and_then(|out| String::from_utf8(out.stdout).ok())
.unwrap_or_else(|| "unknown".to_string()); .unwrap_or_else(|| "unknown".to_string());
let version = format!("{} ({} {})", env!("CARGO_PKG_VERSION"), git_date.trim(), git_hash.trim()); let version = format!(
"{} ({} {})",
env!("CARGO_PKG_VERSION"),
git_date.trim(),
git_hash.trim()
);
println!("cargo:rustc-env=APP_VERSION={}", version); println!("cargo:rustc-env=APP_VERSION={}", version);
} }
stub/src/bin/skeletal_client.rs
+6 -7
View File
@@ -1,14 +1,15 @@
use futures_util::StreamExt;
use reqwest::Client; use reqwest::Client;
use std::env; use std::env;
use std::time::Duration; use std::time::Duration;
use futures_util::StreamExt;
#[tokio::main] #[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> { async fn main() -> Result<(), Box<dyn std::error::Error>> {
tracing_subscriber::fmt::init(); tracing_subscriber::fmt::init();
let target = env::var("MCP_TARGET").unwrap_or_else(|_| "https://127.0.0.1:3000".to_string()); let target = env::var("MCP_TARGET").unwrap_or_else(|_| "https://127.0.0.1:3000".to_string());
let token = env::var("MCP_AUTH_TOKEN").unwrap_or_else(|_| "jP76lUJ5DtFRZmcvXH8LKdCTIkp29eAf".to_string()); let token = env::var("MCP_AUTH_TOKEN")
.unwrap_or_else(|_| "jP76lUJ5DtFRZmcvXH8LKdCTIkp29eAf".to_string());
tracing::info!("Starting skeletal client to {}", target); tracing::info!("Starting skeletal client to {}", target);
@@ -20,10 +21,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
tracing::info!("Connecting to SSE: {}", sse_url); tracing::info!("Connecting to SSE: {}", sse_url);
let res = client.get(&sse_url) let res = client.get(&sse_url).bearer_auth(&token).send().await?;
.bearer_auth(&token)
.send()
.await?;
if !res.status().is_success() { if !res.status().is_success() {
tracing::error!("Failed to connect to SSE: {}", res.status()); tracing::error!("Failed to connect to SSE: {}", res.status());
@@ -78,7 +76,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
tracing::info!("Sending test payload to {}", post_url); tracing::info!("Sending test payload to {}", post_url);
tracing::info!("Payload: {}", payload); tracing::info!("Payload: {}", payload);
let post_res = client.post(&post_url) let post_res = client
.post(&post_url)
.bearer_auth(&token) .bearer_auth(&token)
.header("Content-Type", "application/json") .header("Content-Type", "application/json")
.body(payload.to_string()) .body(payload.to_string())
stub/src/main.rs
+2 -4
View File
@@ -1,6 +1,6 @@
use std::sync::Arc;
use clap::Parser; use clap::Parser;
use futures_util::{SinkExt, StreamExt}; use futures_util::{SinkExt, StreamExt};
use std::sync::Arc;
use tokio::io::AsyncBufReadExt; use tokio::io::AsyncBufReadExt;
use tokio::sync::mpsc; use tokio::sync::mpsc;
@@ -94,7 +94,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
tracing::info!("Attempting to connect to {}", ws_url); tracing::info!("Attempting to connect to {}", ws_url);
use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::client::IntoClientRequest;
let mut request = match ws_url.clone().into_client_request() { let request = match ws_url.clone().into_client_request() {
Ok(req) => req, Ok(req) => req,
Err(e) => { Err(e) => {
tracing::error!("Failed to parse target URL {}: {}", ws_url, e); tracing::error!("Failed to parse target URL {}: {}", ws_url, e);
@@ -165,5 +165,3 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
Ok(()) Ok(())
}) })
} }
stub/tests/e2e.rs
+29 -14
View File
@@ -1,5 +1,5 @@
use serde_json::{json, Value}; use serde_json::{Value, json};
use std::io::{BufRead, BufReader, Read, Write}; use std::io::{BufRead, BufReader, Write};
use std::process::{Command, Stdio}; use std::process::{Command, Stdio};
use std::time::Duration; use std::time::Duration;
@@ -18,10 +18,15 @@ fn read_message(reader: &mut impl BufRead) -> Option<Value> {
serde_json::from_str(line.trim()).ok() serde_json::from_str(line.trim()).ok()
} }
#[tokio::test] #[tokio::test]
async fn test_full_system_e2e_performance() { async fn test_full_system_e2e_performance() {
let temp_dir = std::env::temp_dir().join(format!("mcp_e2e_{}", std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs())); let temp_dir = std::env::temp_dir().join(format!(
"mcp_e2e_{}",
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs()
));
std::fs::create_dir_all(&temp_dir).unwrap(); std::fs::create_dir_all(&temp_dir).unwrap();
let test_port = "3042"; // Use a distinct port let test_port = "3042"; // Use a distinct port
@@ -47,10 +52,13 @@ async fn test_full_system_e2e_performance() {
assert!(stub_exe.exists(), "Stub not found at {:?}", stub_exe); assert!(stub_exe.exists(), "Stub not found at {:?}", stub_exe);
// 1. Start Server // 1. Start Server
let mut server = Command::new(&server_exe).arg("--daemon") let mut server = Command::new(&server_exe)
.env("MCP_PORT", test_port).env("RUST_LOG", "debug") .arg("--daemon")
.env("MCP_PORT", test_port)
.env("RUST_LOG", "debug")
.env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap()) .env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())
.env("MCP_AUTH_TOKEN", test_auth_token).env("RUST_LOG", "debug") .env("MCP_AUTH_TOKEN", test_auth_token)
.env("RUST_LOG", "debug")
.stdout(Stdio::inherit()) .stdout(Stdio::inherit())
.stderr(Stdio::inherit()) .stderr(Stdio::inherit())
.spawn() .spawn()
@@ -63,12 +71,15 @@ async fn test_full_system_e2e_performance() {
.unwrap(); .unwrap();
let mut started = false; let mut started = false;
for _ in 0..30 { for _ in 0..30 {
if let Ok(resp) = client.get(format!("http://127.0.0.1:{}/health", test_port)).send().await { if let Ok(resp) = client
if resp.status().is_success() { .get(format!("http://127.0.0.1:{}/health", test_port))
.send()
.await
&& resp.status().is_success()
{
started = true; started = true;
break; break;
} }
}
tokio::time::sleep(Duration::from_millis(500)).await; tokio::time::sleep(Duration::from_millis(500)).await;
} }
assert!(started, "Server failed to start in time"); assert!(started, "Server failed to start in time");
@@ -78,7 +89,8 @@ async fn test_full_system_e2e_performance() {
.arg("--target") .arg("--target")
.arg(format!("http://127.0.0.1:{}", test_port)) .arg(format!("http://127.0.0.1:{}", test_port))
.env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap()) .env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())
.env("MCP_AUTH_TOKEN", test_auth_token).env("RUST_LOG", "debug") .env("MCP_AUTH_TOKEN", test_auth_token)
.env("RUST_LOG", "debug")
.stdin(Stdio::piped()) .stdin(Stdio::piped())
.stdout(Stdio::piped()) .stdout(Stdio::piped())
.stderr(Stdio::inherit()) .stderr(Stdio::inherit())
@@ -121,7 +133,8 @@ async fn test_full_system_e2e_performance() {
stub_stdin.flush().unwrap(); stub_stdin.flush().unwrap();
} }
let mut resp = read_message(&mut stub_stdout).expect("Failed to read rapid response from stub"); let mut resp =
read_message(&mut stub_stdout).expect("Failed to read rapid response from stub");
while resp.get("id").is_none() || resp["id"].is_null() { while resp.get("id").is_none() || resp["id"].is_null() {
resp = read_message(&mut stub_stdout).expect("Failed to read rapid response from stub"); resp = read_message(&mut stub_stdout).expect("Failed to read rapid response from stub");
} }
@@ -149,9 +162,11 @@ async fn test_full_system_e2e_performance() {
nvim_stdin.flush().unwrap(); nvim_stdin.flush().unwrap();
} }
let mut resp = read_message(&mut nvim_stdout).expect("Failed to read rapid response from win-nvim"); let mut resp =
read_message(&mut nvim_stdout).expect("Failed to read rapid response from win-nvim");
while resp.get("id").is_none() || resp["id"].is_null() { while resp.get("id").is_none() || resp["id"].is_null() {
resp = read_message(&mut nvim_stdout).expect("Failed to read rapid response from win-nvim"); resp = read_message(&mut nvim_stdout)
.expect("Failed to read rapid response from win-nvim");
} }
assert_eq!(resp["id"], i); assert_eq!(resp["id"], i);
} }
stub/tests/negative_scenarios.rs
+25 -5
View File
@@ -10,7 +10,11 @@ fn get_stub_exe() -> std::path::PathBuf {
#[tokio::test] #[tokio::test]
async fn test_stub_connection_refused() { async fn test_stub_connection_refused() {
let _ = std::process::Command::new("cargo").arg("build").arg("--bin").arg("mcp-memory-stub").status(); let _ = std::process::Command::new("cargo")
.arg("build")
.arg("--bin")
.arg("mcp-memory-stub")
.status();
let target = "http://127.0.0.1:49999"; let target = "http://127.0.0.1:49999";
@@ -25,12 +29,20 @@ async fn test_stub_connection_refused() {
let res = tokio::time::timeout(Duration::from_secs(5), child.wait()).await; let res = tokio::time::timeout(Duration::from_secs(5), child.wait()).await;
let elapsed = start.elapsed(); let elapsed = start.elapsed();
assert!(res.is_ok(), "Stub hung on connection refused! Took {:?}", elapsed); assert!(
res.is_ok(),
"Stub hung on connection refused! Took {:?}",
elapsed
);
} }
#[tokio::test] #[tokio::test]
async fn test_stub_handles_eof_cleanly() { async fn test_stub_handles_eof_cleanly() {
let _ = std::process::Command::new("cargo").arg("build").arg("--bin").arg("mcp-memory-stub").status(); let _ = std::process::Command::new("cargo")
.arg("build")
.arg("--bin")
.arg("mcp-memory-stub")
.status();
let target = "http://127.0.0.1:49998"; let target = "http://127.0.0.1:49998";
@@ -58,7 +70,11 @@ async fn test_stub_handles_eof_cleanly() {
#[tokio::test] #[tokio::test]
async fn test_stub_sse_fallback_failure() { async fn test_stub_sse_fallback_failure() {
let _ = std::process::Command::new("cargo").arg("build").arg("--bin").arg("mcp-memory-stub").status(); let _ = std::process::Command::new("cargo")
.arg("build")
.arg("--bin")
.arg("mcp-memory-stub")
.status();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let local_addr = listener.local_addr().unwrap(); let local_addr = listener.local_addr().unwrap();
@@ -84,5 +100,9 @@ async fn test_stub_sse_fallback_failure() {
let res = tokio::time::timeout(Duration::from_secs(5), child.wait()).await; let res = tokio::time::timeout(Duration::from_secs(5), child.wait()).await;
let elapsed = start.elapsed(); let elapsed = start.elapsed();
assert!(res.is_ok(), "Stub hung on fallback failure! Took {:?}", elapsed); assert!(
res.is_ok(),
"Stub hung on fallback failure! Took {:?}",
elapsed
);
} }
win-nvim/build.rs
+8 -3
View File
@@ -2,19 +2,24 @@ use std::process::Command;
fn main() { fn main() {
let git_hash = Command::new("git") let git_hash = Command::new("git")
.args(&["rev-parse", "--short", "HEAD"]) .args(["rev-parse", "--short", "HEAD"])
.output() .output()
.ok() .ok()
.and_then(|out| String::from_utf8(out.stdout).ok()) .and_then(|out| String::from_utf8(out.stdout).ok())
.unwrap_or_else(|| "unknown".to_string()); .unwrap_or_else(|| "unknown".to_string());
let git_date = Command::new("git") let git_date = Command::new("git")
.args(&["log", "-1", "--format=%cd", "--date=format:%Y.%m.%d"]) .args(["log", "-1", "--format=%cd", "--date=format:%Y.%m.%d"])
.output() .output()
.ok() .ok()
.and_then(|out| String::from_utf8(out.stdout).ok()) .and_then(|out| String::from_utf8(out.stdout).ok())
.unwrap_or_else(|| "unknown".to_string()); .unwrap_or_else(|| "unknown".to_string());
let version = format!("{} ({} {})", env!("CARGO_PKG_VERSION"), git_date.trim(), git_hash.trim()); let version = format!(
"{} ({} {})",
env!("CARGO_PKG_VERSION"),
git_date.trim(),
git_hash.trim()
);
println!("cargo:rustc-env=APP_VERSION={}", version); println!("cargo:rustc-env=APP_VERSION={}", version);
} }
win-nvim/tests/integration_test.rs
+7 -2
View File
@@ -112,7 +112,9 @@ fn test_mcp_initialization_and_tools_list() {
assert_eq!(tools_resp["jsonrpc"], "2.0"); assert_eq!(tools_resp["jsonrpc"], "2.0");
assert_eq!(tools_resp["id"], 2); assert_eq!(tools_resp["id"], 2);
let tools = tools_resp["result"]["tools"].as_array().expect("result.tools must be an array"); let tools = tools_resp["result"]["tools"]
.as_array()
.expect("result.tools must be an array");
assert!(!tools.is_empty(), "Server must expose at least one tool"); assert!(!tools.is_empty(), "Server must expose at least one tool");
let has_get_active_buffer = tools.iter().any(|t| t["name"] == "nvim_get_active_buffer"); let has_get_active_buffer = tools.iter().any(|t| t["name"] == "nvim_get_active_buffer");
@@ -137,7 +139,10 @@ fn test_mcp_initialization_and_tools_list() {
assert_eq!(call_resp["jsonrpc"], "2.0"); assert_eq!(call_resp["jsonrpc"], "2.0");
assert_eq!(call_resp["id"], 3); assert_eq!(call_resp["id"], 3);
assert!(call_resp.get("error").is_some(), "Expected an error response since Neovim shouldn't be running"); assert!(
call_resp.get("error").is_some(),
"Expected an error response since Neovim shouldn't be running"
);
assert_eq!(call_resp["error"]["code"], -32603); // Internal Error assert_eq!(call_resp["error"]["code"], -32603); // Internal Error
child.kill().expect("Failed to kill child"); child.kill().expect("Failed to kill child");