From 368f28e593dcb5ab5bca27e94fc86a9d664b5bb9 Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Sat, 19 Sep 2026 07:17:57 +0100 Subject: [PATCH] style: Optimize stub proxy logging and cleanup nvim-core JSON-RPC response boilerplate --- Cargo.lock | 1 + nvim-core/src/lib.rs | 92 +++++++++++--------------------------------- patch_nvim_macro.py | 40 +++++++++++++++++++ patch_stub.py | 15 ++++++++ stub/Cargo.toml | 1 + stub/src/logger.rs | 30 +++++++++++++++ stub/src/main.rs | 23 ++--------- 7 files changed, 113 insertions(+), 89 deletions(-) create mode 100644 patch_nvim_macro.py create mode 100644 patch_stub.py create mode 100644 stub/src/logger.rs diff --git a/Cargo.lock b/Cargo.lock index aa3f8d4..1b85d02 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1376,6 +1376,7 @@ dependencies = [ "dirs 7.0.0", "futures-util", "mcp-stdio", + "regex", "reqwest", "serde_json", "tokio", diff --git a/nvim-core/src/lib.rs b/nvim-core/src/lib.rs index 6f1b1b3..dd2f412 100644 --- a/nvim-core/src/lib.rs +++ b/nvim-core/src/lib.rs @@ -476,6 +476,20 @@ async fn execute_nvim_lua(code: &str) -> Result { } 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!( @@ -649,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, } @@ -666,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, }, @@ -713,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, } @@ -742,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, } @@ -773,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, } @@ -790,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, } diff --git a/patch_nvim_macro.py b/patch_nvim_macro.py new file mode 100644 index 0000000..1da0e4a --- /dev/null +++ b/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") diff --git a/patch_stub.py b/patch_stub.py new file mode 100644 index 0000000..7b3de08 --- /dev/null +++ b/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::\(&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::\(&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") diff --git a/stub/Cargo.toml b/stub/Cargo.toml index e439649..b9e51f3 100644 --- a/stub/Cargo.toml +++ b/stub/Cargo.toml @@ -17,6 +17,7 @@ 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" diff --git a/stub/src/logger.rs b/stub/src/logger.rs new file mode 100644 index 0000000..a9ceb30 --- /dev/null +++ b/stub/src/logger.rs @@ -0,0 +1,30 @@ +use std::sync::LazyLock; +use regex::Regex; + +static ID_REGEX: LazyLock = LazyLock::new(|| Regex::new(r#""id"\s*:\s*([^,}]+)"#).unwrap()); +static METHOD_REGEX: LazyLock = LazyLock::new(|| Regex::new(r#""method"\s*:\s*"([^"]+)""#).unwrap()); +static TOOL_REGEX: LazyLock = LazyLock::new(|| Regex::new(r#""name"\s*:\s*"([^"]+)""#).unwrap()); +static ERROR_REGEX: LazyLock = LazyLock::new(|| Regex::new(r#""error"\s*:\s*\{"#).unwrap()); +static IS_ERROR_REGEX: LazyLock = 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() + } +} diff --git a/stub/src/main.rs b/stub/src/main.rs index 137999d..eba20f3 100644 --- a/stub/src/main.rs +++ b/stub/src/main.rs @@ -14,6 +14,8 @@ struct Cli { +mod logger; + fn init_logging(app_name: &str) -> Option { let mut base_dir = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from(".")); base_dir.push(".gemini/mcp_memory/logs"); @@ -87,19 +89,7 @@ fn main() -> Result<(), Box> { match rx.recv().await { Some(msg) => { drop(rx); - let mut log_prefix = String::from("Message"); - if let Ok(v) = serde_json::from_str::(&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() { @@ -115,12 +105,7 @@ fn main() -> Result<(), Box> { 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::(&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;