Compare commits

...
4 Commits
4 changed files with 81 additions and 5 deletions

No files matched your search

+14 -1
View File
@@ -524,9 +524,18 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
tokio::spawn(async move {
let id = msg.id.clone().unwrap_or(json!(null));
let _start_time = std::time::Instant::now();
let id_clone = id.clone();
let start_time = std::time::Instant::now();
let method_clone = if msg.method == "tools/call" {
let tool_name = msg.params.as_ref().and_then(|p| p.get("name")).and_then(|n| n.as_str()).unwrap_or("unknown");
format!("ToolCall[{}]", tool_name)
} else {
msg.method.clone()
};
match msg.method.as_str() {
"initialize" => {
let init = rmcp::model::InitializeResult::new(
rmcp::model::ServerCapabilities::builder()
@@ -818,10 +827,14 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
}
}
}
let elapsed = start_time.elapsed();
tracing::info!("<<< [Nvim] {} (id: {}) completed in {:?}", method_clone, id_clone, elapsed);
});
}
}
fn init_logging(app_name: &str) -> tracing_appender::non_blocking::WorkerGuard {
let log_dir = dirs::home_dir()
.unwrap_or_default()
+24
View File
@@ -33,10 +33,23 @@ pub struct MemoryHandler {
impl MemoryHandler {
pub async fn handle_request(&self, req: serde_json::Value) -> Option<serde_json::Value> {
let start_time = std::time::Instant::now();
let id = req.get("id").cloned().unwrap_or(serde_json::Value::Null);
let id_clone = id.clone();
let method = req.get("method").and_then(|m| m.as_str()).unwrap_or("");
let tool_name = if method == "tools/call" {
req.get("params").and_then(|p| p.get("name")).and_then(|n| n.as_str()).unwrap_or("unknown")
} else {
""
};
if method == "tools/call" {
tracing::info!(">>> [Server] Handling MCP tool call: {} (id: {})", tool_name, id);
} else {
tracing::debug!(">>> [Server] Handling MCP request method: {}", method);
}
tracing::trace!(">>> [Server] Full MCP Request payload: {}", req.to_string());
let response = match method {
"server/discover" => {
@@ -1648,6 +1661,17 @@ impl MemoryHandler {
}
};
let elapsed = start_time.elapsed();
if method == "tools/call" {
let is_error = response.as_ref().map_or(false, |r| r.get("error").is_some() || r.get("result").and_then(|res| res.get("isError")).and_then(|e| e.as_bool()).unwrap_or(false));
tracing::info!("<<< [Server] MCP tool call {} (id: {}) completed in {:?} [Error: {}]", tool_name, id_clone, elapsed, is_error);
// Broadcast completion latency to the UI Activity Feed
let status_msg = if is_error { "with error" } else { "successfully" };
self.state.broadcast_activity(&format!("Tool {} completed {} in {:?}", tool_name, status_msg, elapsed));
} else {
tracing::debug!("<<< [Server] MCP request method {} completed in {:?}", method, elapsed);
}
tracing::trace!("Returning response from handle_request: {:?}", response);
response
}
+21 -2
View File
@@ -263,6 +263,7 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
"/shutdown",
post(|| async move {
std::thread::spawn(|| {
tracing::info!("Received shutdown request via /shutdown endpoint. Exiting process cleanly.");
std::thread::sleep(std::time::Duration::from_millis(100));
std::process::exit(0);
});
@@ -406,7 +407,16 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
let addr = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
let addr: std::net::SocketAddr = format!("127.0.0.1:{}", addr).parse().unwrap();
let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
let listener = match tokio::net::TcpListener::bind(addr).await {
Ok(l) => l,
Err(e) => {
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_error.log");
let _ = std::fs::write(&log_path, format!("Failed to bind to {}: {}\n", addr, e));
return Ok(());
}
};
if let Err(e) = axum::serve(listener, app.into_make_service()).await {
let log_path = dirs::home_dir()
.unwrap_or_default()
@@ -744,7 +754,16 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
let state = Arc::new(MemoryState {
graph: Store::new("knowledge_graph_master", db.clone()),
base_dir: base.clone(),
search_index: RwLock::new(crate::search::MemoryIndex::new(&base).unwrap()),
search_index: RwLock::new(match crate::search::MemoryIndex::new(&base) {
Ok(idx) => idx,
Err(e) => {
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_error.log");
let _ = std::fs::write(&log_path, format!("Failed to create MemoryIndex: {}\n", e));
std::process::exit(1);
}
}),
ledger: Store::new("audit_ledger", db.clone()),
sticky: Store::new("sticky_notes", db.clone()),
tasks: Store::new("tasks", db.clone()),
+22 -2
View File
@@ -115,8 +115,22 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
match rx.recv().await {
Some(msg) => {
drop(rx);
tracing::info!("Forwarding message to target server (length: {}): {}", msg.len(), if msg.len() > 1000 { format!("{}...", &msg[..1000]) } else { msg.clone() });
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);
}
}
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;
}
@@ -129,7 +143,13 @@ 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 {
tracing::info!("Received message from target server (length: {}): {}", text.len(), if text.len() > 1000 { format!("{}...", &text[..1000]) } else { text.clone() });
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);
}
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;
let mut stdout = tokio::io::stdout();