style: Optimize stub proxy logging and cleanup nvim-core JSON-RPC response boilerplate

Cargo.lock

@@ -1375,6 +1375,8 @@ dependencies = [
  "clap",
  "dirs 7.0.0",
  "futures-util",
+ "mcp-stdio",
+ "regex",
  "reqwest",
  "serde_json",
  "tokio",
@@ -1404,6 +1406,14 @@ dependencies = [
  "tracing-subscriber",
 ]
 
+[[package]]
+name = "mcp-stdio"
+version = "0.1.0"
+dependencies = [
+ "tokio",
+ "tracing",
+]
+
 [[package]]
 name = "measure_time"
 version = "0.9.0"
@@ -1527,6 +1537,7 @@ name = "nvim-core"
 version = "0.1.0"
 dependencies = [
  "dirs 7.0.0",
+ "mcp-stdio",
  "rmcp",
  "rmp-serde",
  "rmpv",

Cargo.toml

@@ -4,5 +4,5 @@ members = [
     "stub",
     "win-nvim",
     "linux-nvim"
-, "nvim-core"]
+, "nvim-core", "mcp-stdio"]
 resolver = "2"

build.ps1

@@ -20,7 +20,12 @@ if (Test-Path $serverExe) {
     & $serverExe --exit 2>$null
 }
 try {
-    Invoke-RestMethod -Uri "https://127.0.0.1:3000/shutdown" -Method Post -SkipCertificateCheck -ErrorAction Stop | Out-Null
+    $tokenPath = "C:\Users\reazul.ashraf\.gemini\mcp_memory\admin.token"
+    $headers = @{}
+    if (Test-Path $tokenPath) {
+        $headers["Authorization"] = "Bearer $(Get-Content -Path $tokenPath -Raw)".Trim()
+    }
+    Invoke-RestMethod -Uri "https://127.0.0.1:3000/shutdown" -Method Post -Headers $headers -SkipCertificateCheck -ErrorAction Stop | Out-Null
 } catch {
     # The response often ends prematurely because the server abruptly kills its own process during shutdown. This is expected.
 }
@@ -63,7 +68,7 @@ function Deploy-If-Needed {
     if ($EnvName -eq 'WSL') {
         # Translate source path to WSL path
         $wslSource = wsl.exe -d Ubuntu -e wslpath -u $Source
-        wsl.exe -d Ubuntu -e bash -c "mv $Dest ${Dest}.`date +%s`.old 2>/dev/null || true; cp $wslSource $Dest && chmod +x $Dest"
+        wsl.exe -d Ubuntu -e bash -c "mv $Dest ${Dest}.\$(date +%s).old 2>/dev/null || true; cp $wslSource $Dest && chmod +x $Dest"
     } else {
         if (Test-Path $Dest) {
             $rand = Get-Random

deploy.ps1

@@ -1,6 +1,11 @@
 try {
     Write-Host "Gracefully shutting down server..."
-    Invoke-RestMethod -Method Post -Uri "http://127.0.0.1:3000/shutdown" -ErrorAction SilentlyContinue
+    $tokenPath = "C:\Users\reazul.ashraf\.gemini\mcp_memory\admin.token"
+    $headers = @{}
+    if (Test-Path $tokenPath) {
+        $headers["Authorization"] = "Bearer $(Get-Content -Path $tokenPath -Raw)".Trim()
+    }
+    Invoke-RestMethod -Method Post -Uri "http://127.0.0.1:3000/shutdown" -Headers $headers -ErrorAction SilentlyContinue | Out-Null
     Start-Sleep -Seconds 2
 } catch {
     Write-Host "Server might already be down or unreachable."

mcp-stdio/Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "mcp-stdio"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+tokio = { version = "1.53.1", features = ["io-util"] }
+tracing = "0.1.44"

mcp-stdio/src/lib.rs

@@ -0,0 +1,40 @@
+use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
+
+/// Reads an MCP (NDJSON or LSP Content-Length prefixed) message from a buffered async reader.
+/// Returns the raw JSON string payload if successful, or None on EOF or error.
+pub async fn read_mcp_message<R: tokio::io::AsyncRead + Unpin>(
+    stdin: &mut BufReader<R>,
+) -> Option<String> {
+    let mut length = 0;
+    loop {
+        let mut line = String::new();
+        if stdin.read_line(&mut line).await.unwrap_or(0) == 0 {
+            return None;
+        }
+
+        if line.starts_with('{') {
+            return Some(line.trim_end().to_string());
+        }
+
+        let line = line.trim_end();
+        if line.is_empty() {
+            break;
+        }
+        
+        let lower_line = line.to_lowercase();
+        if let Some(len_str) = lower_line.strip_prefix("content-length:") {
+            length = len_str.trim().parse().unwrap_or(0);
+        }
+    }
+    
+    if length == 0 {
+        return None;
+    }
+    
+    let mut buffer = vec![0; length];
+    if stdin.read_exact(&mut buffer).await.is_err() {
+        return None;
+    }
+    
+    String::from_utf8(buffer).ok()
+}

nvim-core/Cargo.toml

@@ -14,4 +14,5 @@ tracing-appender = "0.2.5"
 tracing-subscriber = "0.3.23"
 dirs = "7.0.0"
 rmcp = { version = "3.4.0", features = ["server"] }
+mcp-stdio = { version = "0.1.0", path = "../mcp-stdio" }
 

nvim-core/src/lib.rs

@@ -20,47 +20,7 @@ pub struct JsonRpcResponse {
     pub error: Option<Value>,
 }
 
-pub async fn read_message<R: tokio::io::AsyncRead + Unpin>(
-    stdin: &mut BufReader<R>,
-) -> Option<JsonRpcRequest> {
-    let mut length = 0;
-    loop {
-        let mut line = String::new();
-        if stdin.read_line(&mut line).await.unwrap_or(0) == 0 {
-            return None;
-        }
 
-        if line.starts_with('{') {
-            return match serde_json::from_str::<JsonRpcRequest>(line.trim_end()) {
-                Ok(req) => Some(req),
-                Err(e) => {
-                    tracing::error!(
-                        "Failed to parse JSON-RPC request from JSONL: {}. Payload: {}",
-                        e,
-                        line
-                    );
-                    None
-                }
-            };
-        }
-
-        let line = line.trim_end();
-        if line.is_empty() {
-            break;
-        }
-        let lower_line = line.to_lowercase();
-        if let Some(len_str) = lower_line.strip_prefix("content-length:") {
-            length = len_str.trim().parse().unwrap_or(0);
-        }
-    }
-    if length == 0 {
-        return None;
-    }
-    let mut buffer = vec![0; length];
-    stdin.read_exact(&mut buffer).await.unwrap_or(0);
-
-    serde_json::from_slice(&buffer).ok()
-}
 
 pub async fn send_response(response: JsonRpcResponse) {
     let msg = serde_json::to_string(&response).unwrap();
@@ -99,7 +59,7 @@ async fn get_socket_path() -> Result<String, String> {
         std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into());
     let path = format!("{}\\.gemini\\active_nvim.txt", profile);
 
-    if let Ok(content) = std::fs::read_to_string(&path) {
+    if let Ok(content) = tokio::fs::read_to_string(&path).await {
         let p = content.trim().to_string();
         if !p.is_empty() {
             if p.starts_with(r"\\.\pipe\") {
@@ -113,8 +73,8 @@ async fn get_socket_path() -> Result<String, String> {
     }
 
     tracing::warn!("active_nvim.txt missing or invalid, falling back to pipe discovery");
-    if let Ok(dir) = std::fs::read_dir(r"\\.\pipe\") {
-        for entry in dir.flatten() {
+    if let Ok(mut dir) = tokio::fs::read_dir(r"\\.\pipe\").await {
+        while let Ok(Some(entry)) = dir.next_entry().await {
             let name = entry.file_name();
             let name_str = name.to_string_lossy();
             if name_str.starts_with("nvim.") {
@@ -129,7 +89,7 @@ async fn get_socket_path() -> Result<String, String> {
 async fn get_socket_path() -> Result<String, String> {
     if let Ok(home) = std::env::var("HOME") {
         let path = format!("{}/.gemini/active_nvim.txt", home);
-        if let Ok(content) = std::fs::read_to_string(&path) {
+        if let Ok(content) = tokio::fs::read_to_string(&path).await {
             let p = content.trim().to_string();
             if !p.is_empty() {
                 return Ok(p);
@@ -137,12 +97,12 @@ async fn get_socket_path() -> Result<String, String> {
         }
     }
 
-    if let Ok(entries) = std::fs::read_dir("/tmp") {
-        for entry in entries.flatten() {
+    if let Ok(mut entries) = tokio::fs::read_dir("/tmp").await {
+        while let Ok(Some(entry)) = entries.next_entry().await {
             if let Ok(name) = entry.file_name().into_string() {
                 if name.starts_with("nvim.") {
                     let path = entry.path().join("0");
-                    if path.exists() {
+                    if tokio::fs::try_exists(&path).await.unwrap_or(false) {
                         return Ok(path.to_string_lossy().to_string());
                     }
                 }
@@ -151,95 +111,69 @@ async fn get_socket_path() -> Result<String, String> {
     }
     Err("Could not find Neovim socket".to_string())
 }
-#[cfg(windows)]
-async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
-    use tokio::net::windows::named_pipe::ClientOptions;
+use std::sync::LazyLock;
+use std::sync::Arc;
+use tokio::sync::{Mutex, mpsc, oneshot};
+use std::collections::HashMap;
 
-    let msgid = if let rmpv::Value::Array(ref arr) = req {
-        if arr.len() > 1 {
-            arr[1].clone()
-        } else {
-            rmpv::Value::Nil
+
+pub struct NvimRequest {
+    pub msgid_str: String,
+    pub req: rmpv::Value,
+    pub reply: oneshot::Sender<Result<rmpv::Value, String>>,
 }
-    } else {
-        rmpv::Value::Nil
-    };
 
-    tracing::info!("Connecting to neovim pipe");
+static NVIM_CONN: LazyLock<Arc<Mutex<Option<mpsc::Sender<NvimRequest>>>>> = LazyLock::new(|| Arc::new(Mutex::new(None)));
+
+async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
+    let mut conn_lock = NVIM_CONN.lock().await;
+    if let Some(sender) = conn_lock.as_ref() {
+        if !sender.is_closed() {
+            return Ok(sender.clone());
+        }
+    }
+
+    tracing::info!("Establishing new persistent connection to Neovim");
     let socket_path = get_socket_path().await?;
-    let mut client = ClientOptions::new()
-        .open(&socket_path)
-        .map_err(|e| e.to_string())?;
 
-    let mut buf = Vec::new();
-    rmpv::encode::write_value(&mut buf, &req).map_err(|e| e.to_string())?;
-    tracing::info!("Sending RPC request to neovim (msgid: {})", msgid);
-    client.write_all(&buf).await.map_err(|e| e.to_string())?;
-
-    let mut resp_buf = Vec::new();
-    let mut chunk = vec![0u8; 8192];
-    let mut offset = 0;
-
-    loop {
-        let mut cursor = std::io::Cursor::new(&resp_buf[offset..]);
-        match rmpv::decode::read_value(&mut cursor) {
-            Ok(val) => {
-                offset += cursor.position() as usize;
-
-                if let rmpv::Value::Array(ref arr) = val {
-                    if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) && arr[1] == msgid
-                    {
-                        tracing::info!("Received RPC response from neovim (msgid: {})", msgid);
-                        return Ok(val);
-                    }
-                }
-                continue;
-            }
-            Err(_) => {
-                let read_future = client.read(&mut chunk);
-                match tokio::time::timeout(tokio::time::Duration::from_secs(5), read_future).await {
-                    Ok(Ok(n)) => {
-                        if n == 0 {
-                            return Err("Connection closed".into());
-                        }
-                        resp_buf.extend_from_slice(&chunk[..n]);
-                    }
-                    Ok(Err(e)) => return Err(e.to_string()),
-                    Err(_) => {
-                        tracing::error!("Timeout waiting for Neovim response (msgid: {})", msgid);
-                        return Err("Timeout waiting for Neovim response".into());
-                    }
-                }
-            }
-        }
-    }
-}
+    #[cfg(windows)]
+    let stream = {
+        use tokio::net::windows::named_pipe::ClientOptions;
+        ClientOptions::new().open(&socket_path).map_err(|e| e.to_string())?
+    };
 
     #[cfg(unix)]
-async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
+    let stream = {
         use tokio::net::UnixStream;
-
-    let msgid = if let rmpv::Value::Array(ref arr) = req {
-        if arr.len() > 1 {
-            arr[1].clone()
-        } else {
-            rmpv::Value::Nil
-        }
-    } else {
-        rmpv::Value::Nil
+        UnixStream::connect(socket_path).await.map_err(|e| e.to_string())?
     };
 
-    tracing::info!("Connecting to neovim socket");
-    let socket_path = get_socket_path().await?;
-    let mut stream = UnixStream::connect(socket_path)
-        .await
-        .map_err(|e| e.to_string())?;
+    let (mut read_half, mut write_half) = tokio::io::split(stream);
+    let (tx, mut rx) = mpsc::channel::<NvimRequest>(32);
+    let pending_requests: Arc<Mutex<HashMap<String, oneshot::Sender<Result<rmpv::Value, String>>>>> = Arc::new(Mutex::new(HashMap::new()));
 
+    // Write task
+    let pending_clone = Arc::clone(&pending_requests);
+    tokio::spawn(async move {
+        while let Some(req) = rx.recv().await {
             let mut buf = Vec::new();
-    rmpv::encode::write_value(&mut buf, &req).map_err(|e| e.to_string())?;
-    tracing::info!("Sending RPC request to neovim (msgid: {})", msgid);
-    stream.write_all(&buf).await.map_err(|e| e.to_string())?;
+            if let Err(e) = rmpv::encode::write_value(&mut buf, &req.req) {
+                let _ = req.reply.send(Err(e.to_string()));
+                continue;
+            }
             
+            pending_clone.lock().await.insert(req.msgid_str.clone(), req.reply);
+            
+            if write_half.write_all(&buf).await.is_err() {
+                tracing::error!("Failed to write to Neovim socket");
+                break;
+            }
+        }
+    });
+
+    // Read task
+    let pending_clone2 = Arc::clone(&pending_requests);
+    tokio::spawn(async move {
         let mut resp_buf = Vec::new();
         let mut chunk = vec![0u8; 8192];
         let mut offset = 0;
@@ -251,33 +185,81 @@ async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
                     offset += cursor.position() as usize;
 
                     if let rmpv::Value::Array(ref arr) = val {
-                    if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) && arr[1] == msgid
-                    {
-                        tracing::info!("Received RPC response from neovim (msgid: {})", msgid);
-                        return Ok(val);
+                        if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) {
+                            let msgid = &arr[1];
+                            let msgid_str = format!("{:?}", msgid);
+                            
+                            if let Some(reply_sender) = pending_clone2.lock().await.remove(&msgid_str) {
+                                let _ = reply_sender.send(Ok(val));
                             }
                         }
+                    }
+                    // Trim buffer if it gets too large
+                    if offset > 1024 * 1024 {
+                        resp_buf.drain(..offset);
+                        offset = 0;
+                    }
                     continue;
                 }
                 Err(_) => {
-                let read_future = stream.read(&mut chunk);
-                match tokio::time::timeout(tokio::time::Duration::from_secs(5), read_future).await {
-                    Ok(Ok(n)) => {
-                        if n == 0 {
-                            return Err("Connection closed".into());
+                    if offset > 0 {
+                        resp_buf.drain(..offset);
+                        offset = 0;
                     }
+                    
+                    let read_future = read_half.read(&mut chunk);
+                    match tokio::time::timeout(tokio::time::Duration::from_secs(60), read_future).await {
+                        Ok(Ok(n)) if n > 0 => {
                             resp_buf.extend_from_slice(&chunk[..n]);
                         }
-                    Ok(Err(e)) => return Err(e.to_string()),
-                    Err(_) => {
-                        tracing::error!("Timeout waiting for Neovim response (msgid: {})", msgid);
-                        return Err("Timeout waiting for Neovim response".into());
+                        _ => {
+                            tracing::error!("Neovim socket read loop closed or timeout");
+                            break;
                         }
                     }
                 }
             }
         }
+        
+        // Cleanup pending requests on disconnect
+        let mut pending = pending_clone2.lock().await;
+        for (_, sender) in pending.drain() {
+            let _ = sender.send(Err("Connection closed".to_string()));
         }
+    });
+
+    *conn_lock = Some(tx.clone());
+    Ok(tx)
+}
+
+async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
+    let msgid = if let rmpv::Value::Array(ref arr) = req {
+        if arr.len() > 1 {
+            arr[1].clone()
+        } else {
+            rmpv::Value::Nil
+        }
+    } else {
+        rmpv::Value::Nil
+    };
+    
+    let msgid_str = format!("{:?}", msgid);
+    let tx = get_nvim_connection().await?;
+    let (reply_tx, reply_rx) = oneshot::channel();
+    
+    tx.send(NvimRequest {
+        msgid_str,
+        req,
+        reply: reply_tx,
+    }).await.map_err(|_| "Failed to send request to Neovim connection manager")?;
+    
+    match tokio::time::timeout(tokio::time::Duration::from_secs(10), reply_rx).await {
+        Ok(Ok(res)) => res,
+        Ok(Err(_)) => Err("Response channel dropped".to_string()),
+        Err(_) => Err("Timeout waiting for Neovim response".to_string()),
+    }
+}
+
 async fn send_nvim_command(cmd: &str) -> Result<(), String> {
     use rmpv::Value as RmpValue;
     let req = RmpValue::Array(vec![
@@ -494,6 +476,20 @@ async fn execute_nvim_lua(code: &str) -> Result<String, String> {
     }
     Err("Invalid response".to_string())
 }
+
+macro_rules! send_text_result {
+    ($id:expr, $text:expr) => {
+        send_response(JsonRpcResponse {
+            jsonrpc: "2.0".to_string(),
+            id: $id,
+            result: Some(json!({
+                "content": [{"type": "text", "text": $text}]
+            })),
+            error: None,
+        }).await;
+    };
+}
+
 pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
     if std::env::args().any(|arg| arg == "--version") {
         println!(
@@ -508,17 +504,25 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
     tracing::info!("{} MCP server started", app_name);
     let mut stdin = tokio::io::BufReader::new(tokio::io::stdin());
     loop {
-        let msg = match read_message(&mut stdin).await {
-            Some(m) => {
-                tracing::info!("Received message method: {}", m.method);
-                m
-            }
+        let raw_msg = match mcp_stdio::read_mcp_message(&mut stdin).await {
+            Some(m) => m,
             None => {
                 tracing::info!("Stdin closed, exiting loop");
                 break;
             }
         };
 
+        let msg = match serde_json::from_str::<JsonRpcRequest>(&raw_msg) {
+            Ok(m) => {
+                tracing::info!("Received message method: {}", m.method);
+                m
+            },
+            Err(e) => {
+                tracing::error!("Failed to parse JSON-RPC request from JSONL: {}. Payload: {}", e, raw_msg);
+                continue;
+            }
+        };
+
         let app_name = app_name.to_string();
         let app_version = app_version.to_string();
 
@@ -659,14 +663,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
                                 let cmd = format!("e {} | {} | normal! zz", escaped_file, line);
                                 match send_nvim_command(&cmd).await {
                                     Ok(_) => {
-                                        send_response(JsonRpcResponse {
-                                            jsonrpc: "2.0".to_string(),
-                                            id,
-                                            result: Some(json!({
-                                                "content": [{"type": "text", "text": format!("Successfully jumped to {} line {}", file, line)}]
-                                            })),
-                                            error: None,
-                                        }).await;
+                                        send_text_result!(id.clone(), format!("Successfully jumped to {} line {}", file, line));
                                     }
                                     Err(e) => send_error(id, -32603, &e).await,
                                 }
@@ -676,43 +673,19 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
                         }
                         "nvim_get_active_buffer" => match get_nvim_active_buffer().await {
                             Ok(content) => {
-                                send_response(JsonRpcResponse {
-                                    jsonrpc: "2.0".to_string(),
-                                    id,
-                                    result: Some(json!({
-                                        "content": [{"type": "text", "text": content}]
-                                    })),
-                                    error: None,
-                                })
-                                .await;
+                                send_text_result!(id.clone(), content);
                             }
                             Err(e) => send_error(id, -32603, &e).await,
                         },
                         "nvim_get_cursor" => match get_nvim_cursor().await {
                             Ok(content) => {
-                                send_response(JsonRpcResponse {
-                                    jsonrpc: "2.0".to_string(),
-                                    id,
-                                    result: Some(json!({
-                                        "content": [{"type": "text", "text": content}]
-                                    })),
-                                    error: None,
-                                })
-                                .await;
+                                send_text_result!(id.clone(), content);
                             }
                             Err(e) => send_error(id, -32603, &e).await,
                         },
                         "nvim_get_visual_selection" => match get_nvim_visual_selection().await {
                             Ok(content) => {
-                                send_response(JsonRpcResponse {
-                                    jsonrpc: "2.0".to_string(),
-                                    id,
-                                    result: Some(json!({
-                                        "content": [{"type": "text", "text": content}]
-                                    })),
-                                    error: None,
-                                })
-                                .await;
+                                send_text_result!(id.clone(), content);
                             }
                             Err(e) => send_error(id, -32603, &e).await,
                         },
@@ -723,14 +696,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
                             ) {
                                 match set_nvim_diagnostics(line, message).await {
                                     Ok(_) => {
-                                        send_response(JsonRpcResponse {
-                                            jsonrpc: "2.0".to_string(),
-                                            id,
-                                            result: Some(json!({
-                                                "content": [{"type": "text", "text": format!("Successfully set diagnostic on line {}", line)}]
-                                            })),
-                                            error: None,
-                                        }).await;
+                                        send_text_result!(id.clone(), format!("Successfully set diagnostic on line {}", line));
                                     }
                                     Err(e) => send_error(id, -32603, &e).await,
                                 }
@@ -752,15 +718,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
                             "#;
                             match execute_nvim_lua(lua_code).await {
                                 Ok(result) => {
-                                    send_response(JsonRpcResponse {
-                                        jsonrpc: "2.0".to_string(),
-                                        id,
-                                        result: Some(json!({
-                                            "content": [{"type": "text", "text": result}]
-                                        })),
-                                        error: None,
-                                    })
-                                    .await;
+                                    send_text_result!(id.clone(), result);
                                 }
                                 Err(e) => send_error(id, -32603, &e).await,
                             }
@@ -783,15 +741,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
                             "#;
                             match execute_nvim_lua(lua_code).await {
                                 Ok(result) => {
-                                    send_response(JsonRpcResponse {
-                                        jsonrpc: "2.0".to_string(),
-                                        id,
-                                        result: Some(json!({
-                                            "content": [{"type": "text", "text": result}]
-                                        })),
-                                        error: None,
-                                    })
-                                    .await;
+                                    send_text_result!(id.clone(), result);
                                 }
                                 Err(e) => send_error(id, -32603, &e).await,
                             }
@@ -800,15 +750,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
                             if let Some(code) = args.get("code").and_then(|v| v.as_str()) {
                                 match execute_nvim_lua(code).await {
                                     Ok(result) => {
-                                        send_response(JsonRpcResponse {
-                                            jsonrpc: "2.0".to_string(),
-                                            id,
-                                            result: Some(json!({
-                                                "content": [{"type": "text", "text": result}]
-                                            })),
-                                            error: None,
-                                        })
-                                        .await;
+                                        send_text_result!(id.clone(), result);
                                     }
                                     Err(e) => send_error(id, -32603, &e).await,
                                 }
@@ -915,3 +857,4 @@ mod tests {
         assert!(req.is_none());
     }
 }
+

patch_nvim.py

@@ -0,0 +1,168 @@
+import re
+import sys
+
+with open('C:\\Users\\reazul.ashraf\\workspace\\rust\\mcp-memory\\nvim-core\\src\\lib.rs', 'r', encoding='utf-8') as f:
+    content = f.read()
+
+start_pattern = r'#\[cfg\(windows\)\]\nasync fn call_nvim'
+end_pattern = r'async fn send_nvim_command'
+
+start_idx = re.search(start_pattern, content).start()
+end_idx = re.search(end_pattern, content).start()
+
+new_code = """use std::sync::LazyLock;
+use std::sync::Arc;
+use tokio::sync::{Mutex, mpsc, oneshot};
+use std::collections::HashMap;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+pub struct NvimRequest {
+    pub msgid_str: String,
+    pub req: rmpv::Value,
+    pub reply: oneshot::Sender<Result<rmpv::Value, String>>,
+}
+
+static NVIM_CONN: LazyLock<Arc<Mutex<Option<mpsc::Sender<NvimRequest>>>>> = LazyLock::new(|| Arc::new(Mutex::new(None)));
+
+async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
+    let mut conn_lock = NVIM_CONN.lock().await;
+    if let Some(sender) = conn_lock.as_ref() {
+        if !sender.is_closed() {
+            return Ok(sender.clone());
+        }
+    }
+
+    tracing::info!("Establishing new persistent connection to Neovim");
+    let socket_path = get_socket_path().await?;
+
+    #[cfg(windows)]
+    let stream = {
+        use tokio::net::windows::named_pipe::ClientOptions;
+        ClientOptions::new().open(&socket_path).map_err(|e| e.to_string())?
+    };
+
+    #[cfg(unix)]
+    let stream = {
+        use tokio::net::UnixStream;
+        UnixStream::connect(socket_path).await.map_err(|e| e.to_string())?
+    };
+
+    let (mut read_half, mut write_half) = tokio::io::split(stream);
+    let (tx, mut rx) = mpsc::channel::<NvimRequest>(32);
+    let pending_requests: Arc<Mutex<HashMap<String, oneshot::Sender<Result<rmpv::Value, String>>>>> = Arc::new(Mutex::new(HashMap::new()));
+
+    // Write task
+    let pending_clone = Arc::clone(&pending_requests);
+    tokio::spawn(async move {
+        while let Some(req) = rx.recv().await {
+            let mut buf = Vec::new();
+            if let Err(e) = rmpv::encode::write_value(&mut buf, &req.req) {
+                let _ = req.reply.send(Err(e.to_string()));
+                continue;
+            }
+            
+            pending_clone.lock().await.insert(req.msgid_str.clone(), req.reply);
+            
+            if write_half.write_all(&buf).await.is_err() {
+                tracing::error!("Failed to write to Neovim socket");
+                break;
+            }
+        }
+    });
+
+    // Read task
+    let pending_clone2 = Arc::clone(&pending_requests);
+    tokio::spawn(async move {
+        let mut resp_buf = Vec::new();
+        let mut chunk = vec![0u8; 8192];
+        let mut offset = 0;
+
+        loop {
+            let mut cursor = std::io::Cursor::new(&resp_buf[offset..]);
+            match rmpv::decode::read_value(&mut cursor) {
+                Ok(val) => {
+                    offset += cursor.position() as usize;
+
+                    if let rmpv::Value::Array(ref arr) = val {
+                        if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) {
+                            let msgid = &arr[1];
+                            let msgid_str = format!("{:?}", msgid);
+                            
+                            if let Some(reply_sender) = pending_clone2.lock().await.remove(&msgid_str) {
+                                let _ = reply_sender.send(Ok(val));
+                            }
+                        }
+                    }
+                    // Trim buffer if it gets too large
+                    if offset > 1024 * 1024 {
+                        resp_buf.drain(..offset);
+                        offset = 0;
+                    }
+                    continue;
+                }
+                Err(_) => {
+                    if offset > 0 {
+                        resp_buf.drain(..offset);
+                        offset = 0;
+                    }
+                    
+                    let read_future = read_half.read(&mut chunk);
+                    match tokio::time::timeout(tokio::time::Duration::from_secs(60), read_future).await {
+                        Ok(Ok(n)) if n > 0 => {
+                            resp_buf.extend_from_slice(&chunk[..n]);
+                        }
+                        _ => {
+                            tracing::error!("Neovim socket read loop closed or timeout");
+                            break;
+                        }
+                    }
+                }
+            }
+        }
+        
+        // Cleanup pending requests on disconnect
+        let mut pending = pending_clone2.lock().await;
+        for (_, sender) in pending.drain() {
+            let _ = sender.send(Err("Connection closed".to_string()));
+        }
+    });
+
+    *conn_lock = Some(tx.clone());
+    Ok(tx)
+}
+
+async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
+    let msgid = if let rmpv::Value::Array(ref arr) = req {
+        if arr.len() > 1 {
+            arr[1].clone()
+        } else {
+            rmpv::Value::Nil
+        }
+    } else {
+        rmpv::Value::Nil
+    };
+    
+    let msgid_str = format!("{:?}", msgid);
+    let tx = get_nvim_connection().await?;
+    let (reply_tx, reply_rx) = oneshot::channel();
+    
+    tx.send(NvimRequest {
+        msgid_str,
+        req,
+        reply: reply_tx,
+    }).await.map_err(|_| "Failed to send request to Neovim connection manager")?;
+    
+    match tokio::time::timeout(tokio::time::Duration::from_secs(10), reply_rx).await {
+        Ok(Ok(res)) => res,
+        Ok(Err(_)) => Err("Response channel dropped".to_string()),
+        Err(_) => Err("Timeout waiting for Neovim response".to_string()),
+    }
+}
+
+"""
+
+new_content = content[:start_idx] + new_code + content[end_idx:]
+with open('C:\\Users\\reazul.ashraf\\workspace\\rust\\mcp-memory\\nvim-core\\src\\lib.rs', 'w', encoding='utf-8') as f:
+    f.write(new_content)
+
+print("Patched!")

patch_nvim_macro.py

@@ -0,0 +1,40 @@
+import re
+import sys
+
+with open('C:\\Users\\reazul.ashraf\\workspace\\rust\\mcp-memory\\nvim-core\\src\\lib.rs', 'r', encoding='utf-8') as f:
+    content = f.read()
+
+macro_def = """
+macro_rules! send_text_result {
+    ($id:expr, $text:expr) => {
+        send_response(JsonRpcResponse {
+            jsonrpc: "2.0".to_string(),
+            id: $id,
+            result: Some(json!({
+                "content": [{"type": "text", "text": $text}]
+            })),
+            error: None,
+        }).await;
+    };
+}
+"""
+
+if "macro_rules! send_text_result" not in content:
+    idx = content.find("pub async fn run_mcp_loop")
+    content = content[:idx] + macro_def + "\n" + content[idx:]
+
+# Find tools/call block
+tools_call_idx = content.find('"tools/call" => {')
+if tools_call_idx != -1:
+    block = content[tools_call_idx:]
+    
+    # Replace the verbose blocks
+    pattern = r'send_response\(JsonRpcResponse \{\s*jsonrpc: "2\.0"\.to_string\(\),\s*id,\s*result: Some\(json!\(\{\s*"content": \[\{"type": "text", "text": (.*?)\}\]\s*\}\)\),\s*error: None,\s*\}\)\s*\.await;'
+    
+    new_block = re.sub(pattern, r'send_text_result!(id.clone(), \1);', block)
+    
+    content = content[:tools_call_idx] + new_block
+
+with open('C:\\Users\\reazul.ashraf\\workspace\\rust\\mcp-memory\\nvim-core\\src\\lib.rs', 'w', encoding='utf-8') as f:
+    f.write(content)
+print("Patched nvim macro")

patch_stub.py

@@ -0,0 +1,15 @@
+import re
+import sys
+
+with open('C:\\Users\\reazul.ashraf\\workspace\\rust\\mcp-memory\\stub\\src\\main.rs', 'r', encoding='utf-8') as f:
+    content = f.read()
+
+start1 = r'let mut log_prefix = String::from\("Message"\);\n\s*if let Ok\(v\) = serde_json::from_str::<serde_json::Value>\(&msg\) \{[\s\S]*?log_prefix = format!\("Response id=\{\}", res_id\);\n\s*\}\n\s*\}'
+content = re.sub(start1, r'let log_prefix = logger::extract_log_prefix(&msg, false);', content)
+
+start2 = r'let mut log_prefix = String::from\("Message"\);\n\s*if let Ok\(v\) = serde_json::from_str::<serde_json::Value>\(&text\) \{[\s\S]*?log_prefix = format!\("Response id=\{\} \[Error: \{\}\]", id, is_error\);\n\s*\}'
+content = re.sub(start2, r'let log_prefix = logger::extract_log_prefix(&text, true);', content)
+
+with open('C:\\Users\\reazul.ashraf\\workspace\\rust\\mcp-memory\\stub\\src\\main.rs', 'w', encoding='utf-8') as f:
+    f.write(content)
+print("Patched stub")

server/src/main.rs

@@ -261,14 +261,33 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
             .route("/gate/set", post(gate_set_handler))
             .route(
                 "/shutdown",
-                post(|| async move {
+                post(
+                    |headers: axum::http::HeaderMap, State(state): State<Arc<AppState>>| async move {
+                        let token_path = state.handler.state.base_dir.join("admin.token");
+                        let expected_token = std::fs::read_to_string(&token_path)
+                            .unwrap_or_default()
+                            .trim()
+                            .to_string();
+
+                        let auth_header = headers
+                            .get(axum::http::header::AUTHORIZATION)
+                            .and_then(|h| h.to_str().ok())
+                            .unwrap_or_default();
+
+                        if expected_token.is_empty() || auth_header != format!("Bearer {}", expected_token) {
+                            return (axum::http::StatusCode::UNAUTHORIZED, "Unauthorized").into_response();
+                        }
+
                         std::thread::spawn(|| {
-                        tracing::info!("Received shutdown request via /shutdown endpoint. Exiting process cleanly.");
+                            tracing::info!(
+                                "Received shutdown request via /shutdown endpoint. Exiting process cleanly."
+                            );
                             std::thread::sleep(std::time::Duration::from_millis(100));
                             std::process::exit(0);
                         });
-                    "Shutting down..."
-                }),
+                        (axum::http::StatusCode::OK, "Shutting down...").into_response()
+                    },
+                ),
             )
             .route(
                 "/",
@@ -648,26 +667,38 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let _guard = init_logging("mcp-memory-server");
     let cli = Cli::parse();
 
+    let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
+        dirs::home_dir()
+            .map(|mut h| {
+                h.push(".gemini/mcp_memory");
+                h.to_string_lossy().into_owned()
+            })
+            .unwrap_or_else(|| ".gemini/mcp_memory".into())
+    });
+    let base = PathBuf::from(base_dir);
+
     if cli.exit {
         let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
-        let _ = std::process::Command::new("curl")
-            .arg("-k")
-            .arg("-X")
-            .arg("POST")
-            .arg(format!("https://127.0.0.1:{}/shutdown", port))
-            .output();
+        let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default();
+        let mut cmd = std::process::Command::new("curl");
+        cmd.arg("-k").arg("-X").arg("POST");
+        if !token.is_empty() {
+            cmd.arg("-H").arg(format!("Authorization: Bearer {}", token.trim()));
+        }
+        let _ = cmd.arg(format!("https://127.0.0.1:{}/shutdown", port)).output();
         println!("Sent shutdown request to server.");
         return Ok(());
     }
 
     if cli.restart {
         let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
-        let _ = std::process::Command::new("curl")
-            .arg("-k")
-            .arg("-X")
-            .arg("POST")
-            .arg(format!("https://127.0.0.1:{}/shutdown", port))
-            .output();
+        let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default();
+        let mut cmd = std::process::Command::new("curl");
+        cmd.arg("-k").arg("-X").arg("POST");
+        if !token.is_empty() {
+            cmd.arg("-H").arg(format!("Authorization: Bearer {}", token.trim()));
+        }
+        let _ = cmd.arg(format!("https://127.0.0.1:{}/shutdown", port)).output();
         println!("Sent shutdown request to existing server. Waiting for it to exit...");
         std::thread::sleep(std::time::Duration::from_millis(1500));
         return Ok(());
@@ -691,18 +722,14 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         }
     }
 
-    let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
-        dirs::home_dir()
-            .map(|mut h| {
-                h.push(".gemini/mcp_memory");
-                h.to_string_lossy().into_owned()
-            })
-            .unwrap_or_else(|| ".gemini/mcp_memory".into())
-    });
-    let base = PathBuf::from(base_dir);
     fs::create_dir_all(&base).expect("Failed to create store dir");
 
+    // Generate token
+    let admin_token = uuid::Uuid::new_v4().to_string();
+    std::fs::write(base.join("admin.token"), &admin_token).expect("Failed to write admin token");
+
     let redb_path = base.join("mcp_store.redb");
+
     let db = Arc::new(redb::Database::create(&redb_path).unwrap());
 
     // Ensure table exists and migrate old JSON files

stub/Cargo.toml

@@ -16,6 +16,8 @@ tracing = "0.1.44"
 tracing-subscriber = "0.3.23"
 dirs = "7.0.0"
 serde_json = "1.0.151"
+mcp-stdio = { version = "0.1.0", path = "../mcp-stdio" }
+regex = "1.13.1"
 
 
 

stub/src/logger.rs

@@ -0,0 +1,30 @@
+use std::sync::LazyLock;
+use regex::Regex;
+
+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 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());
+
+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");
+    
+    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("");
+    
+    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");
+        format!("ToolCall[{}] id={}", tool, id)
+    } else if !method.is_empty() {
+        format!("Request[{}] id={}", method, id)
+    } else if id != "null" {
+        format!("Response id={}", id)
+    } else {
+        "Message".to_string()
+    }
+}

stub/src/main.rs

@@ -12,40 +12,9 @@ struct Cli {
     target: String,
 }
 
-async fn read_mcp_message(stdin: &mut tokio::io::BufReader<tokio::io::Stdin>) -> Option<String> {
-    use tokio::io::AsyncReadExt;
-    let mut length = 0;
-    loop {
-        let mut line = String::new();
-        let bytes_read = stdin.read_line(&mut line).await.unwrap_or(0);
-        if bytes_read == 0 {
-            tracing::info!("stdin EOF reached");
-            return None;
-        }
-        tracing::info!("Read {} bytes from stdin: {:?}", bytes_read, line);
 
-        if line.starts_with('{') {
-            return Some(line.trim_end().to_string());
-        }
 
-        let line = line.trim_end();
-        if line.is_empty() {
-            break;
-        }
-        let lower_line = line.to_lowercase();
-        if let Some(len_str) = lower_line.strip_prefix("content-length:") {
-            length = len_str.trim().parse().unwrap_or(0);
-        }
-    }
-    if length == 0 {
-        return None;
-    }
-    let mut buffer = vec![0; length];
-    if stdin.read_exact(&mut buffer).await.is_err() {
-        return None;
-    }
-    String::from_utf8(buffer).ok()
-}
+mod logger;
 
 fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
     let mut base_dir = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from("."));
@@ -74,13 +43,18 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
         tokio::spawn(async move {
             let mut stdin = tokio::io::BufReader::new(tokio::io::stdin());
-            while let Some(msg) = read_mcp_message(&mut stdin).await {
+            while let Some(msg) = mcp_stdio::read_mcp_message(&mut stdin).await {
                 let _ = msg_tx.send(msg).await;
             }
             let _ = shutdown_tx.send(()).await;
         });
 
-        let target_url = cli.target;
+        let target_url = if cli.target != "http://localhost:3000" {
+            cli.target
+        } else {
+            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 = format!("{}/ws?client=proxy", ws_url);
         let msg_rx = Arc::new(tokio::sync::Mutex::new(msg_rx));
@@ -115,19 +89,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                             match rx.recv().await {
                                 Some(msg) => {
                                     drop(rx);
-                                    let mut log_prefix = String::from("Message");
-                                    if let Ok(v) = serde_json::from_str::<serde_json::Value>(&msg) {
-                                        let method = v.get("method").and_then(|m| m.as_str()).unwrap_or("");
-                                        let id = v.get("id").unwrap_or(&serde_json::Value::Null);
-                                        if method == "tools/call" {
-                                            let tool = v.get("params").and_then(|p| p.get("name")).and_then(|n| n.as_str()).unwrap_or("unknown");
-                                            log_prefix = format!("ToolCall[{}] id={}", tool, id);
-                                        } else if !method.is_empty() {
-                                            log_prefix = format!("Request[{}] id={}", method, id);
-                                        } else if let Some(res_id) = v.get("id") {
-                                            log_prefix = format!("Response id={}", res_id);
-                                        }
-                                    }
+                                    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() {
 
@@ -143,12 +105,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                     let mut recv_task = tokio::spawn(async move {
                         while let Some(Ok(msg)) = read.next().await {
                             if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
-                                let mut log_prefix = String::from("Message");
-                                if let Ok(v) = serde_json::from_str::<serde_json::Value>(&text) {
-                                    let id = v.get("id").unwrap_or(&serde_json::Value::Null);
-                                    let is_error = v.get("error").is_some() || v.get("result").and_then(|res| res.get("isError")).and_then(|e| e.as_bool()).unwrap_or(false);
-                                    log_prefix = format!("Response id={} [Error: {}]", id, is_error);
-                                }
+                                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() });
                                 let payload = format!("{}\n", text);
                                 use tokio::io::AsyncWriteExt;
@@ -185,3 +142,4 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         Ok(())
     })
 }
+