diff --git a/crates/mcp/src/backend.rs b/crates/mcp/src/backend.rs index f3686545..cc43b4ee 100644 --- a/crates/mcp/src/backend.rs +++ b/crates/mcp/src/backend.rs @@ -55,16 +55,24 @@ impl Remote { Err(last) } - fn roundtrip(&mut self, line: &str) -> std::io::Result { - if self.conn.is_none() { + /// Write one request line. A failure here means the app never got a complete line, so retrying is safe. + /// A connection the app has already closed (it restarted since the last call) is replaced before writing. + fn send(&mut self, line: &str) -> std::io::Result<()> { + if self.conn.as_ref().is_none_or(|(_, writer)| is_closed(writer)) { self.reconnect()?; } - let Some((reader, writer)) = self.conn.as_mut() else { + let Some((_, writer)) = self.conn.as_mut() else { + return Err(std::io::Error::new(std::io::ErrorKind::NotConnected, "not connected")); + }; + writer.write_all(format!("{line}\n").as_bytes())?; + writer.flush() + } + + /// Read the reply to a request that was already sent. + fn receive(&mut self) -> std::io::Result { + let Some((reader, _)) = self.conn.as_mut() else { return Err(std::io::Error::new(std::io::ErrorKind::NotConnected, "not connected")); }; - writer.write_all(line.as_bytes())?; - writer.write_all(b"\n")?; - writer.flush()?; let mut reply = String::new(); if reader.read_line(&mut reply)? == 0 { return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "control channel closed")); @@ -73,22 +81,39 @@ impl Remote { } } +/// True when the peer has closed or reset `stream`, or sent bytes nobody asked for. Between calls nothing +/// should be waiting, so only "would block" means the connection is still good. +fn is_closed(stream: &TcpStream) -> bool { + if stream.set_nonblocking(true).is_err() { + return true; + } + let mut byte = [0u8; 1]; + let open = matches!(stream.peek(&mut byte), Err(e) if e.kind() == std::io::ErrorKind::WouldBlock); + stream.set_nonblocking(false).is_err() || !open +} + impl Backend for Remote { fn call(&mut self, method: &str, params: Value) -> Result { let id = self.next_id; self.next_id += 1; let line = json!({"id": id, "method": method, "params": params}).to_string(); - // One retry with a fresh connection (the app may have restarted). - let reply = match self.roundtrip(&line) { - Ok(r) => r, - Err(_) => { + // One retry with a fresh connection (the app may have restarted), but only while the request + // was not sent. Once it is out, a lost reply leaves the outcome unknown and a resend could + // run an insert/delete/undo twice, so report that instead. + if self.send(&line).is_err() { + self.conn = None; + self.send(&line).map_err(|e| { self.conn = None; - self.roundtrip(&line).map_err(|e| { - self.conn = None; - format!("DeckCraft app at {} is not reachable: {e}", self.addr) - })? - } - }; + format!("DeckCraft app at {} is not reachable: {e}", self.addr) + })?; + } + let reply = self.receive().map_err(|e| { + self.conn = None; + format!( + "lost the reply from the DeckCraft app at {} ({e}); the command may or may not have run, check the document before retrying", + self.addr + ) + })?; let v: Value = serde_json::from_str(reply.trim()).map_err(|e| format!("bad reply from app: {e}"))?; if v.get("ok").and_then(Value::as_bool) == Some(true) { Ok(v.get("result").cloned().unwrap_or(Value::Null)) @@ -105,3 +130,75 @@ impl Backend for Remote { format!("connected to the DeckCraft app at {}", self.addr) } } + +#[cfg(test)] +mod tests { + use super::*; + use std::net::TcpListener; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[test] + fn lost_reply_is_not_resent() { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let addr = listener.local_addr().expect("addr").to_string(); + let executed = Arc::new(AtomicUsize::new(0)); + let count = Arc::clone(&executed); + std::thread::spawn(move || { + for stream in listener.incoming().flatten() { + let mut reader = BufReader::new(&stream); + let mut line = String::new(); + if reader.read_line(&mut line).unwrap_or(0) == 0 { + continue; + } + let n = count.fetch_add(1, Ordering::SeqCst) + 1; + if n > 1 { + let mut w = &stream; + let _ = w.write_all(b"{\"id\":2,\"ok\":true,\"result\":null}\n"); + } + // The first request is executed, then the connection drops without a reply. + } + }); + let mut remote = Remote::connect(&addr).expect("connect"); + let err = remote.call("engine.execute", json!({})).expect_err("lost reply must be an error"); + assert!(err.contains("may or may not have run"), "{err}"); + assert_eq!(executed.load(Ordering::SeqCst), 1, "the command must not be replayed"); + // The next call opens a fresh connection and works. + assert!(remote.call("document.inspect", json!({})).is_ok()); + assert_eq!(executed.load(Ordering::SeqCst), 2); + } + + #[test] + fn connection_closed_before_the_call_is_retried() { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let addr = listener.local_addr().expect("addr").to_string(); + let (closed_tx, closed_rx) = std::sync::mpsc::channel(); + let executed = Arc::new(AtomicUsize::new(0)); + let count = Arc::clone(&executed); + std::thread::spawn(move || { + let mut first = true; + for stream in listener.incoming().flatten() { + if first { + // The app restarts: the first connection closes before any request arrives. + first = false; + drop(stream); + let _ = closed_tx.send(()); + continue; + } + let mut reader = BufReader::new(&stream); + let mut line = String::new(); + if reader.read_line(&mut line).unwrap_or(0) == 0 { + continue; + } + count.fetch_add(1, Ordering::SeqCst); + let mut w = &stream; + let _ = w.write_all(b"{\"id\":1,\"ok\":true,\"result\":null}\n"); + } + }); + let mut remote = Remote::connect(&addr).expect("connect"); + closed_rx.recv().expect("server closed the first connection"); + std::thread::sleep(Duration::from_millis(200)); + assert!(remote.call("engine.execute", json!({})).is_ok(), "a request that never left must be sent on a fresh connection"); + assert_eq!(executed.load(Ordering::SeqCst), 1); + } +}