Fix MCP hangs and add logging

stub/Cargo.toml

@@ -10,3 +10,7 @@ tokio = { version = "1.53.1", features = ["full"] }
 tokio-util = { version = "0.7.19", features = ["io"] }
 futures-util = "0.3.34"
 tokio-tungstenite = "0.21.0"
+tracing-appender = "0.2.5"
+tracing = "0.1.44"
+tracing-subscriber = "0.3.23"
+dirs = "7.0.0"

stub/src/main.rs

@@ -15,7 +15,58 @@ struct Cli {
     wake_cmd: Option<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();
+        if stdin.read_line(&mut line).await.unwrap_or(0) == 0 {
+            return None;
+        }
+        let line = line.trim_end();
+        if line.is_empty() {
+            break;
+        }
+        if let Some(len_str) = line.strip_prefix("Content-Length: ") {
+            length = len_str.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()
+}
+
+fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
+    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().to_string()
+            })
+            .unwrap_or_else(|| ".gemini/mcp_memory".into())
+    });
+    let log_dir = std::path::PathBuf::from(base_dir).join("logs");
+    std::fs::create_dir_all(&log_dir).unwrap_or_default();
+    
+    let file_appender = tracing_appender::rolling::daily(log_dir, format!("{}.log", app_name));
+    let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
+    
+    let _ = tracing_subscriber::fmt()
+        .with_writer(non_blocking)
+        .with_ansi(false)
+        .with_max_level(tracing::Level::INFO)
+        .try_init();
+    
+    Some(guard)
+}
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let _guard = init_logging("stub");
     let cli = Cli::parse();
     let rt = tokio::runtime::Runtime::new()?;
     rt.block_on(async {
@@ -24,13 +75,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
         tokio::spawn(async move {
             let mut stdin = tokio::io::BufReader::new(tokio::io::stdin());
-            let mut buffer = String::new();
-            while let Ok(bytes) = stdin.read_line(&mut buffer).await {
-                if bytes == 0 {
-                    break;
-                }
-                let _ = msg_tx.send(buffer.clone()).await;
-                buffer.clear();
+            while let Some(msg) = read_mcp_message(&mut stdin).await {
+                let _ = msg_tx.send(msg).await;
             }
             let _ = shutdown_tx.send(()).await;
         });
@@ -43,11 +89,14 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
         loop {
             if shutdown_rx.try_recv().is_ok() {
+                tracing::info!("Stub shutdown requested");
                 break;
             }
 
+            tracing::info!("Attempting to connect to {}", ws_url);
             match tokio_tungstenite::connect_async(&ws_url).await {
                 Ok((ws_stream, _)) => {
+                    tracing::info!("Successfully connected to target server");
                     let (mut write, mut read) = ws_stream.split();
 
                     let msg_rx_clone = Arc::clone(&msg_rx);
@@ -57,7 +106,9 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                             match rx.recv().await {
                                 Some(msg) => {
                                     drop(rx);
+                                    tracing::info!("Forwarding message to target server");
                                     if write.send(tokio_tungstenite::tungstenite::Message::Text(msg)).await.is_err() {
+                                        tracing::error!("Failed to write to websocket");
                                         break;
                                     }
                                 },
@@ -69,26 +120,37 @@ 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 {
-                                println!("{}", text);
+                                tracing::info!("Received message from target server, proxying to stdout");
+                                let payload = format!("Content-Length: {}\r\n\r\n{}", text.len(), text);
+                                use std::io::Write;
+                                let mut stdout = std::io::stdout();
+                                let _ = stdout.write_all(payload.as_bytes());
+                                let _ = stdout.flush();
                             }
                         }
+                        tracing::error!("Websocket read loop exited");
                     });
 
                     tokio::select! {
                         _ = shutdown_rx.recv() => {
+                            tracing::info!("Shutdown received while connected");
                             return Ok(()); // Stdin closed, exit entirely
                         }
                         _ = &mut send_task => {
+                            tracing::error!("Send task exited");
                             recv_task.abort();
                         }
                         _ = &mut recv_task => {
+                            tracing::error!("Recv task exited");
                             send_task.abort();
                         }
                     }
                     tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
                 }
-                Err(_) => {
+                Err(e) => {
+                    tracing::error!("Failed to connect to target server: {}", e);
                     if let Some(ref cmd) = wake_cmd {
+                        tracing::info!("Executing wake command: {}", cmd);
                         let parts: Vec<&str> = cmd.split_whitespace().collect();
                         if !parts.is_empty() {
                             let _ = std::process::Command::new(parts[0])