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

refactor: apply zero-unwrap policy and optimize locks in store.rs and handlers

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-21 11:34:21 +01:00
1 parent 9f24e66d88
commit 8afbf97b11
33 files changed
+3127 -3207
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

stub/build.rs
+1 -5
View File
@@ -15,10 +15,6 @@ fn main() {
.and_then(|out| String::from_utf8(out.stdout).ok())
.unwrap_or_else(|| "unknown".to_string());
let version = format!(
"{} ({})",
git_date.trim(),
git_hash.trim()
);
let version = format!("{} ({})", git_date.trim(), git_hash.trim());
println!("cargo:rustc-env=APP_VERSION={}", version);
}
stub/src/logger.rs
+25 -10
View File
@@ -1,24 +1,39 @@
use std::sync::LazyLock;
use regex::Regex;
use std::sync::LazyLock;
static ID_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""id"\s*:\s*([^,}]+)"#).unwrap());
static METHOD_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""method"\s*:\s*"([^"]+)""#).unwrap());
static TOOL_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""name"\s*:\s*"([^"]+)""#).unwrap());
static METHOD_REGEX: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r#""method"\s*:\s*"([^"]+)""#).unwrap());
static TOOL_REGEX: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r#""name"\s*:\s*"([^"]+)""#).unwrap());
static ERROR_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""error"\s*:\s*\{"#).unwrap());
static IS_ERROR_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#""isError"\s*:\s*true"#).unwrap());
static IS_ERROR_REGEX: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r#""isError"\s*:\s*true"#).unwrap());
pub fn extract_log_prefix(json_str: &str, is_response: bool) -> String {
let id = ID_REGEX.captures(json_str).and_then(|c| c.get(1)).map(|m| m.as_str()).unwrap_or("null");
let id = ID_REGEX
.captures(json_str)
.and_then(|c| c.get(1))
.map(|m| m.as_str())
.unwrap_or("null");
if is_response {
let is_error = ERROR_REGEX.is_match(json_str) || IS_ERROR_REGEX.is_match(json_str);
return format!("Response id={} [Error: {}]", id, is_error);
}
let method = METHOD_REGEX.captures(json_str).and_then(|c| c.get(1)).map(|m| m.as_str()).unwrap_or("");
let method = METHOD_REGEX
.captures(json_str)
.and_then(|c| c.get(1))
.map(|m| m.as_str())
.unwrap_or("");
if method == "tools/call" {
let tool = TOOL_REGEX.captures(json_str).and_then(|c| c.get(1)).map(|m| m.as_str()).unwrap_or("unknown");
let tool = TOOL_REGEX
.captures(json_str)
.and_then(|c| c.get(1))
.map(|m| m.as_str())
.unwrap_or("unknown");
format!("ToolCall[{}] id={}", tool, id)
} else if !method.is_empty() {
format!("Request[{}] id={}", method, id)
stub/src/main.rs
+31 -10
View File
@@ -10,8 +10,6 @@ struct Cli {
target: String,
}
mod logger;
fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
@@ -53,7 +51,9 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
format!("http://127.0.0.1:{}", port)
};
let ws_url = target_url.replace("http://", "ws://").replace("https://", "wss://");
let ws_url = target_url
.replace("http://", "ws://")
.replace("https://", "wss://");
let ws_url = format!("{}/ws?client=proxy", ws_url);
loop {
@@ -63,7 +63,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
}
tracing::info!("Attempting to connect to {}", ws_url);
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
let request = match ws_url.clone().into_client_request() {
Ok(req) => req,
@@ -83,8 +83,21 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
let mut send_task = tokio::spawn(async move {
while let Ok(msg) = rx.recv().await {
let log_prefix = logger::extract_log_prefix(&msg, false);
tracing::info!(">>> [Stub] Forwarding {} to server (length: {}): {}", log_prefix, msg.len(), if msg.len() > 1000 { format!("{}...", &msg[..1000]) } else { msg.clone() });
if write.send(tokio_tungstenite::tungstenite::Message::Text(msg)).await.is_err() {
tracing::info!(
">>> [Stub] Forwarding {} to server (length: {}): {}",
log_prefix,
msg.len(),
if msg.len() > 1000 {
format!("{}...", &msg[..1000])
} else {
msg.clone()
}
);
if write
.send(tokio_tungstenite::tungstenite::Message::Text(msg))
.await
.is_err()
{
tracing::error!("Failed to write to websocket");
break;
}
@@ -95,7 +108,16 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
while let Some(Ok(msg)) = read.next().await {
if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
let log_prefix = logger::extract_log_prefix(&text, true);
tracing::info!("<<< [Stub] Received {} from server (length: {}): {}", log_prefix, text.len(), if text.len() > 1000 { format!("{}...", &text[..1000]) } else { text.clone() });
tracing::info!(
"<<< [Stub] Received {} from server (length: {}): {}",
log_prefix,
text.len(),
if text.len() > 1000 {
format!("{}...", &text[..1000])
} else {
text.clone()
}
);
use tokio::io::AsyncWriteExt;
let mut stdout = tokio::io::stdout();
let _ = stdout.write_all(text.as_bytes()).await;
@@ -107,14 +129,14 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
});
tokio::select! {
_ = shutdown_rx.recv() => {
_ = shutdown_rx.recv() => {
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
tracing::info!("Shutdown received while connected");
break;
}
_ = &mut send_task => {
tracing::error!("Send task exited");
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
recv_task.abort();
break;
}
@@ -134,4 +156,3 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
Ok(())
})
}
stub/tests/e2e.rs
+34 -28
View File
@@ -61,17 +61,19 @@ async fn test_full_system_e2e_performance() {
assert!(stub_exe.exists(), "Stub not found at {:?}", stub_exe);
// 1. Start Server
let _server = ChildGuard(Command::new(&server_exe)
.arg("--daemon")
.env("MCP_PORT", test_port)
.env("RUST_LOG", "debug")
.env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())
.env("MCP_AUTH_TOKEN", test_auth_token)
.env("RUST_LOG", "debug")
.stdout(Stdio::inherit())
.stderr(Stdio::inherit())
.spawn()
.expect("Failed to start server"));
let _server = ChildGuard(
Command::new(&server_exe)
.arg("--daemon")
.env("MCP_PORT", test_port)
.env("RUST_LOG", "debug")
.env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())
.env("MCP_AUTH_TOKEN", test_auth_token)
.env("RUST_LOG", "debug")
.stdout(Stdio::inherit())
.stderr(Stdio::inherit())
.spawn()
.expect("Failed to start server"),
);
// Give server time to generate TLS cert and start
let client = reqwest::Client::builder()
@@ -94,28 +96,32 @@ async fn test_full_system_e2e_performance() {
assert!(started, "Server failed to start in time");
// 2. Start Stub
let mut stub = ChildGuard(Command::new(&stub_exe)
.arg("--target")
.arg(format!("http://127.0.0.1:{}", test_port))
.env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())
.env("MCP_AUTH_TOKEN", test_auth_token)
.env("RUST_LOG", "debug")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.spawn()
.expect("Failed to start stub"));
let mut stub = ChildGuard(
Command::new(&stub_exe)
.arg("--target")
.arg(format!("http://127.0.0.1:{}", test_port))
.env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())
.env("MCP_AUTH_TOKEN", test_auth_token)
.env("RUST_LOG", "debug")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.spawn()
.expect("Failed to start stub"),
);
let mut stub_stdin = stub.0.stdin.take().unwrap();
let mut stub_stdout = BufReader::new(stub.0.stdout.take().unwrap());
// 3. Start Nvim Bridge
let mut nvim = ChildGuard(Command::new(&nvim_exe)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.spawn()
.expect("Failed to start nvim bridge"));
let mut nvim = ChildGuard(
Command::new(&nvim_exe)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.spawn()
.expect("Failed to start nvim bridge"),
);
let mut nvim_stdin = nvim.0.stdin.take().unwrap();
let mut nvim_stdout = BufReader::new(nvim.0.stdout.take().unwrap());