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

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

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-19 07:17:57 +01:00
1 parent 6011da95f6
commit 368f28e593
7 files changed
+113 -89
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
+1
View File
@@ -1376,6 +1376,7 @@ dependencies = [
"dirs 7.0.0",
"futures-util",
"mcp-stdio",
"regex",
"reqwest",
"serde_json",
"tokio",
nvim-core/src/lib.rs
+22 -70
View File
@@ -476,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!(
@@ -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,
}
patch_nvim_macro.py
+40
View File
@@ -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
+15
View File
@@ -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")
stub/Cargo.toml
+1
View File
@@ -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"
stub/src/logger.rs
+30
View File
@@ -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
+4 -19
View File
@@ -14,6 +14,8 @@ struct Cli {
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("."));
base_dir.push(".gemini/mcp_memory/logs");
@@ -87,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() {
@@ -115,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;