From 707aa5aa9282ed9d53323bda7b07401f0b7ac685 Mon Sep 17 00:00:00 2001 From: euisuh <55032583+euisuh@users.noreply.github.com> Date: Wed, 14 Oct 2026 17:02:08 +0900 Subject: [PATCH] fix(mcp): stop stdio cleanly after output failure Wake the serving loop after protocol output fails and preserve the actual writer error. Use one process-owned stdin worker so the CLI can finish normal backend cleanup without waiting for client EOF, and terminate the output worker even when a progress hook is retained. Closes #1527 Co-Authored-By: OpenAI Codex --- apps/lightcraft-cli/src/main.rs | 5 +- apps/lightcraft-cli/tests/mcp_shutdown.rs | 38 +++++ crates/mcp/src/lib.rs | 3 +- crates/mcp/src/progress.rs | 167 ++++++++++++++------ crates/mcp/src/progress_tests.rs | 177 ++++++++++++++++++++++ crates/mcp/src/server.rs | 4 +- docs/mcp.md | 10 ++ 7 files changed, 356 insertions(+), 48 deletions(-) create mode 100644 apps/lightcraft-cli/tests/mcp_shutdown.rs create mode 100644 crates/mcp/src/progress_tests.rs diff --git a/apps/lightcraft-cli/src/main.rs b/apps/lightcraft-cli/src/main.rs index b6cad57e0..3f34485ba 100644 --- a/apps/lightcraft-cli/src/main.rs +++ b/apps/lightcraft-cli/src/main.rs @@ -17,7 +17,7 @@ mod alloc_release; -use std::io::{BufReader, Write}; +use std::io::Write; use std::path::Path; use std::process::ExitCode; @@ -364,8 +364,7 @@ fn mcp(args: &[String]) -> Result<(), String> { }; eprintln!("lightcraft-cli mcp: serving MCP on stdio ({})", backend.describe()); let mut server = Server::new(backend).with_command_tools(!compact); - let stdin = std::io::stdin(); - server.serve(BufReader::new(stdin), std::io::stdout()).map_err(|e| e.to_string()) + server.serve_stdio(std::io::stdout()).map_err(|e| e.to_string()) } fn merge(args: &[String]) -> Result<(), String> { diff --git a/apps/lightcraft-cli/tests/mcp_shutdown.rs b/apps/lightcraft-cli/tests/mcp_shutdown.rs new file mode 100644 index 000000000..29d1ace93 --- /dev/null +++ b/apps/lightcraft-cli/tests/mcp_shutdown.rs @@ -0,0 +1,38 @@ +//! A broken protocol output must not require the client to close its input pipe. + +use std::io::{Read, Write}; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; + +#[test] +fn broken_stdout_exits_with_stdin_still_open() { + let mut child = Command::new(env!("CARGO_BIN_EXE_lightcraft-cli")) + .arg("mcp") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let mut input = child.stdin.take().unwrap(); + drop(child.stdout.take()); + input.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\"}\n{\"jsonrpc\":").unwrap(); + input.flush().unwrap(); + let deadline = Instant::now() + Duration::from_secs(10); + let status = loop { + if let Some(status) = child.try_wait().unwrap() { + break status; + } + if Instant::now() >= deadline { + child.kill().unwrap(); + child.wait().unwrap(); + panic!("MCP did not exit after stdout failed while stdin stayed open"); + } + std::thread::sleep(Duration::from_millis(10)); + }; + let mut stderr = String::new(); + child.stderr.take().unwrap().read_to_string(&mut stderr).unwrap(); + assert!(stderr.contains("serving MCP on stdio"), "CLI must reach the MCP transport: {stderr}"); + assert!(!status.success(), "broken protocol output must be reported: {stderr}"); + assert!(stderr.contains("lightcraft-cli:"), "CLI must report the transport error: {stderr}"); + drop(input); +} diff --git a/crates/mcp/src/lib.rs b/crates/mcp/src/lib.rs index f95dd7a53..f632f145b 100644 --- a/crates/mcp/src/lib.rs +++ b/crates/mcp/src/lib.rs @@ -11,7 +11,8 @@ //! control-channel method names itself (rendering with the engine's pipeline and encoding with //! `lightcraft-codecs`), so agents can develop photos and look at the result without a window. //! -//! Entry points: [`Server::serve`] (stdio loop) and [`Server::handle_line`] (one message). +//! Entry points: [`Server::serve_stdio`] (process stdin), [`Server::serve`] (caller-managed +//! streams), and [`Server::handle_line`] (one message). #![forbid(unsafe_code)] #![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] diff --git a/crates/mcp/src/progress.rs b/crates/mcp/src/progress.rs index 2043dfd91..3f6a8deec 100644 --- a/crates/mcp/src/progress.rs +++ b/crates/mcp/src/progress.rs @@ -2,8 +2,8 @@ //! serving thread; only stdin and stdout move to transport threads. Other calls wait in order. use std::collections::VecDeque; -use std::io::{BufRead, Write}; -use std::sync::atomic::Ordering; +use std::io::{BufRead, BufReader, Write}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc::{Receiver, Sender, TryRecvError}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -13,12 +13,17 @@ use serde_json::{Value, json}; use super::{Server, response}; use crate::backend::ProgressHook; -type Input = Option>; +enum Input { + Line(std::io::Result), + Closed, + OutputFailed, +} #[derive(Clone)] pub(super) struct Wire { inbox: Arc>>, - out: Sender, + out: Sender>, + failed: Arc, deferred: Arc>>, } @@ -26,56 +31,114 @@ fn lock(m: &Mutex) -> std::sync::MutexGuard<'_, T> { m.lock().unwrap_or_else(std::sync::PoisonError::into_inner) } +fn read_input(input: impl BufRead, tx: Sender) { + let read = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + for line in input.lines() { + let failed = line.is_err(); + if tx.send(Input::Line(line)).is_err() || failed { + return; + } + } + let _ = tx.send(Input::Closed); + })); + if read.is_err() { + let _ = tx.send(Input::Line(Err(std::io::Error::other("MCP input thread panicked")))); + } +} + impl Server { + /// Serve process stdin once, writing replies to `output` (normally stdout). + /// + /// The caller must finish the process after this returns. On output failure, the single + /// stdin worker can remain blocked until process exit; it owns no backend or borrowed data. + /// Further calls are refused, including after EOF, so blocked readers cannot accumulate. + pub fn serve_stdio(&mut self, output: impl Write + Send) -> std::io::Result<()> { + static STARTED: AtomicBool = AtomicBool::new(false); + if STARTED.swap(true, Ordering::AcqRel) { + return Err(std::io::Error::new(std::io::ErrorKind::AlreadyExists, "MCP stdin already claimed")); + } + let (in_tx, in_rx) = std::sync::mpsc::channel(); + let reader_tx = in_tx.clone(); + let reader = match std::thread::Builder::new().name("mcp-input".into()).spawn(move || read_input(BufReader::new(std::io::stdin()), reader_tx)) + { + Ok(reader) => reader, + Err(error) => { + STARTED.store(false, Ordering::Release); + return Err(error); + } + }; + let served = self.serve_messages(in_rx, in_tx, output); + let read = if reader.is_finished() { reader.join().map_err(|_| std::io::Error::other("MCP input thread panicked")) } else { Ok(()) }; + served.and(read) + } + pub(super) fn serve_wire(&mut self, input: impl BufRead + Send, output: impl Write + Send) -> std::io::Result<()> { std::thread::scope(|scope| { - let (in_tx, in_rx) = std::sync::mpsc::channel::(); - let (out_tx, out_rx) = std::sync::mpsc::channel::(); - let wire = Wire { inbox: Arc::new(Mutex::new(in_rx)), out: out_tx, deferred: Arc::new(Mutex::new(VecDeque::new())) }; + let (in_tx, in_rx) = std::sync::mpsc::channel(); + let reader_tx = in_tx.clone(); + let reader = std::thread::Builder::new().name("mcp-input".into()).spawn_scoped(scope, move || read_input(input, reader_tx))?; + let served = self.serve_messages(in_rx, in_tx, output); + let read = reader.join().map_err(|_| std::io::Error::other("MCP input thread panicked")); + served.and(read) + }) + } + + fn serve_messages(&mut self, in_rx: Receiver, in_tx: Sender, output: impl Write + Send) -> std::io::Result<()> { + std::thread::scope(|scope| { + let (out_tx, out_rx) = std::sync::mpsc::channel::>(); + let failed = Arc::new(AtomicBool::new(false)); + let wire = + Wire { inbox: Arc::new(Mutex::new(in_rx)), out: out_tx, failed: failed.clone(), deferred: Arc::new(Mutex::new(VecDeque::new())) }; let writer = std::thread::Builder::new().name("mcp-output".into()).spawn_scoped(scope, move || -> std::io::Result<()> { let mut output = output; - for line in out_rx { - output.write_all(line.as_bytes())?; - output.write_all(b"\n")?; - output.flush()?; - } - Ok(()) - })?; - let reader = match std::thread::Builder::new().name("mcp-input".into()).spawn_scoped(scope, move || { - for line in input.lines() { - let failed = line.is_err(); - if in_tx.send(Some(line)).is_err() || failed { - return; + let written = (|| -> std::io::Result<()> { + while let Ok(Some(line)) = out_rx.recv() { + output.write_all(line.as_bytes())?; + output.write_all(b"\n")?; + output.flush()?; } + Ok(()) + })(); + if written.is_err() { + failed.store(true, Ordering::Release); + let _ = in_tx.send(Input::OutputFailed); } - let _ = in_tx.send(None); - }) { - Ok(reader) => reader, - Err(error) => { - drop(wire); - let _ = writer.join(); - return Err(error); - } - }; + written + })?; self.wire = Some(wire.clone()); let served = (|| -> std::io::Result<()> { loop { + if wire.failed.load(Ordering::Acquire) { + return Ok(()); + } let queued = lock(&wire.deferred).pop_front(); let line = match queued { Some(line) => line, - None => lock(&wire.inbox).recv().ok().flatten(), + None => { + if wire.failed.load(Ordering::Acquire) { + return Ok(()); + } + lock(&wire.inbox).recv().map_err(|_| std::io::Error::other("MCP input disconnected without EOF"))? + } }; - let Some(line) = line else { return Ok(()) }; - if let Some(reply) = self.handle_line(&line?) { - wire.out.send(reply).map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "MCP output closed"))?; + let line = match line { + Input::Line(line) => line?, + Input::Closed | Input::OutputFailed => return Ok(()), + }; + if wire.failed.load(Ordering::Acquire) { + return Ok(()); + } + if let Some(reply) = self.handle_line(&line) { + wire.out.send(Some(reply)).map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "MCP output closed"))?; } } })(); self.wire = None; + // A backend that panicked while clearing its progress hook may still hold a sender. + let _ = wire.out.send(None); drop(wire); let written = writer.join().map_err(|_| std::io::Error::other("MCP output thread panicked")).and_then(|r| r); - let read = reader.join().map_err(|_| std::io::Error::other("MCP input thread panicked")); - served.and(written).and(read) + written.and(served) }) } @@ -91,16 +154,29 @@ impl Server { let mut reported: Option = None; Some(Box::new(move |done, total, name| { while !closed { + if wire.failed.load(Ordering::Acquire) { + suppress.store(true, Ordering::Relaxed); + return false; + } let next = lock(&wire.inbox).try_recv(); let line = match next { - Ok(Some(Ok(line))) => line, - Ok(Some(Err(error))) => { - lock(&wire.deferred).push_back(Some(Err(error))); + Ok(Input::Line(Ok(line))) => line, + Ok(Input::Line(Err(error))) => { + lock(&wire.deferred).push_back(Input::Line(Err(error))); closed = true; break; } - Ok(None) | Err(TryRecvError::Disconnected) => { - lock(&wire.deferred).push_back(None); + Ok(Input::OutputFailed) => { + suppress.store(true, Ordering::Relaxed); + return false; + } + Ok(Input::Closed) => { + lock(&wire.deferred).push_back(Input::Closed); + closed = true; + break; + } + Err(TryRecvError::Disconnected) => { + lock(&wire.deferred).push_back(Input::Line(Err(std::io::Error::other("MCP input disconnected without EOF")))); closed = true; break; } @@ -115,12 +191,13 @@ impl Server { if msg.get("id").is_some_and(|i| i.is_string() || i.is_number()) && msg.get("params").is_none_or(|p| p.is_null() || p.is_object()) => { - let _ = wire.out.send(response(msg["id"].clone(), json!({})).to_string()); + let _ = wire.out.send(Some(response(msg["id"].clone(), json!({})).to_string())); } - _ => lock(&wire.deferred).push_back(Some(Ok(line))), + _ => lock(&wire.deferred).push_back(Input::Line(Ok(line))), } } - if suppress.load(Ordering::Relaxed) { + if suppress.load(Ordering::Relaxed) || wire.failed.load(Ordering::Acquire) { + suppress.store(true, Ordering::Relaxed); return false; } if let Some(token) = &token @@ -137,13 +214,17 @@ impl Server { }; let notification = json!({"jsonrpc":"2.0","method":"notifications/progress", "params":{"progressToken":token,"progress":done,"total":total,"message":message}}); - let _ = wire.out.send(notification.to_string()); + let _ = wire.out.send(Some(notification.to_string())); } true })) } } +#[cfg(test)] +#[path = "progress_tests.rs"] +mod shutdown_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/crates/mcp/src/progress_tests.rs b/crates/mcp/src/progress_tests.rs new file mode 100644 index 000000000..fb7859dc5 --- /dev/null +++ b/crates/mcp/src/progress_tests.rs @@ -0,0 +1,177 @@ +use super::*; +use crate::Backend; +use std::io::{BufReader, Cursor, Read}; +use std::process::{Command, Stdio}; + +struct ExportUntilCancelled { + hook: Option, + retain_hook: bool, +} + +impl Backend for ExportUntilCancelled { + fn call(&mut self, _: &str, _: Value) -> Result { + if self.retain_hook { + return Ok(json!({})); + } + let hook = self.hook.as_mut().expect("export installs a progress hook"); + while hook(0, 2, "photo.png") { + std::thread::sleep(Duration::from_millis(1)); + } + Ok(json!({"cancelled": true})) + } + + fn has_ui(&self) -> bool { + false + } + + fn describe(&self) -> String { + "test".into() + } + + fn set_progress(&mut self, hook: Option) { + assert!(!self.retain_hook || hook.is_some(), "synthetic clear-hook panic"); + self.hook = hook; + } +} + +struct FailedOutput { + flush: bool, + initialized: bool, +} + +impl Write for FailedOutput { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + if !self.initialized || self.flush { Ok(bytes.len()) } else { Err(std::io::Error::other("synthetic output failure")) } + } + + fn flush(&mut self) -> std::io::Result<()> { + if self.initialized { + Err(std::io::Error::other("synthetic output failure")) + } else { + self.initialized = true; + Ok(()) + } + } +} + +// The child owns the blocked read; a regression cannot strand the test runner in a thread join. +#[test] +#[ignore = "subprocess helper for the bounded shutdown regression"] +fn output_failure_child() { + let mode = std::env::var("LIGHTCRAFT_MCP_SHUTDOWN_CASE").unwrap(); + let mut server = Server::new(Box::new(ExportUntilCancelled { hook: None, retain_hook: mode == "retained-hook" })); + if mode == "retained-hook" { + let request = json!({"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"export","arguments":{"dir":"unused"}}}); + let mut output = Vec::new(); + server.serve(Cursor::new(format!("{request}\n")), &mut output).unwrap(); + let reply: Value = serde_json::from_slice(&output).unwrap(); + assert_eq!(reply["result"]["isError"], true); + assert!(reply.to_string().contains("synthetic clear-hook panic")); + } else { + let error = server.serve_stdio(FailedOutput { flush: mode.ends_with("flush"), initialized: false }).unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::Other); + assert_eq!(error.to_string(), "synthetic output failure"); + assert!(server.wire.is_none()); + let mut another = Server::new(Box::new(ExportUntilCancelled { hook: None, retain_hook: false })); + let error = another.serve_stdio(std::io::sink()).unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::AlreadyExists); + } + eprintln!("MCP shutdown verified"); +} + +#[test] +fn idle_write_failure_does_not_wait_for_input_eof() { + run_shutdown_case("idle-write"); +} + +#[test] +fn idle_flush_failure_does_not_wait_for_input_eof() { + run_shutdown_case("idle-flush"); +} + +#[test] +fn progress_write_failure_does_not_wait_for_input_eof() { + run_shutdown_case("progress-write"); +} + +#[test] +fn progress_flush_failure_does_not_wait_for_input_eof() { + run_shutdown_case("progress-flush"); +} + +fn run_shutdown_case(mode: &str) { + let mut child = Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "server::progress::shutdown_tests::output_failure_child", "--ignored", "--nocapture"]) + .env("LIGHTCRAFT_MCP_SHUTDOWN_CASE", mode) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let mut input = child.stdin.take().unwrap(); + if mode != "retained-hook" { + let request = if mode.starts_with("progress") { + json!({"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"export","arguments":{"dir":"unused"},"_meta":{"progressToken":"photos"}}}) + } else { + json!({"jsonrpc":"2.0","id":2,"method":"ping"}) + }; + let initialize = json!({"jsonrpc":"2.0","id":0,"method":"initialize"}); + let requests = format!("{initialize}\n{request}\n{{\"jsonrpc\":"); + input.write_all(requests.as_bytes()).unwrap(); + input.flush().unwrap(); + } + let deadline = Instant::now() + Duration::from_secs(5); + let status = loop { + if let Some(status) = child.try_wait().unwrap() { + break status; + } + if Instant::now() >= deadline { + child.kill().unwrap(); + child.wait().unwrap(); + panic!("{mode}: MCP shutdown waited for input EOF"); + } + std::thread::sleep(Duration::from_millis(10)); + }; + let mut stderr = String::new(); + child.stderr.take().unwrap().read_to_string(&mut stderr).unwrap(); + assert!(status.success(), "{mode}: {status}: {stderr}"); + assert!(stderr.contains("MCP shutdown verified"), "child did not exercise shutdown: {stderr}"); + drop(input); +} + +#[test] +fn retained_progress_hook_does_not_keep_output_alive() { + run_shutdown_case("retained-hook"); +} + +#[test] +fn input_errors_and_eof_keep_their_meaning() { + let mut server = Server::new(Box::new(ExportUntilCancelled { hook: None, retain_hook: false })); + let error = server.serve(Cursor::new(vec![0xff]), std::io::sink()).unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::InvalidData); + assert!(server.wire.is_none()); + server.serve(Cursor::new(b""), std::io::sink()).unwrap(); +} + +#[test] +fn reader_panic_is_an_input_error() { + struct PanickingRead; + impl Read for PanickingRead { + fn read(&mut self, _: &mut [u8]) -> std::io::Result { + panic!("synthetic input panic"); + } + } + let mut server = Server::new(Box::new(ExportUntilCancelled { hook: None, retain_hook: false })); + let error = server.serve(BufReader::new(PanickingRead), std::io::sink()).unwrap_err(); + assert_eq!(error.to_string(), "MCP input thread panicked"); +} + +#[test] +fn output_error_takes_precedence_over_input_error() { + let mut server = Server::new(Box::new(ExportUntilCancelled { hook: None, retain_hook: false })); + let mut input = b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\"}\n".to_vec(); + input.push(0xff); + let error = server.serve(Cursor::new(input), FailedOutput { flush: false, initialized: true }).unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::Other); + assert_eq!(error.to_string(), "synthetic output failure"); +} diff --git a/crates/mcp/src/server.rs b/crates/mcp/src/server.rs index 4565db76e..de086ca70 100644 --- a/crates/mcp/src/server.rs +++ b/crates/mcp/src/server.rs @@ -84,7 +84,9 @@ impl Server { } /// Serve newline-delimited JSON-RPC until `input` closes. Logs go to stderr only (stdout is - /// the protocol stream). + /// the protocol stream). The scoped input worker must reach EOF or an input error before + /// this returns, including on output failure: callers must unblock their own input. + /// For a process serving stdin, use [`Self::serve_stdio`] instead. pub fn serve(&mut self, input: impl BufRead + Send, output: impl Write + Send) -> std::io::Result<()> { self.serve_wire(input, output) } diff --git a/docs/mcp.md b/docs/mcp.md index f0cc9a1d1..ec29a75b7 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -314,5 +314,15 @@ being processed finishes first: completed photos remain, each written with the e file writer, and there are no partial files to delete. Unrelated outputs and earlier exports are untouched. A failed export is an `isError` tool result. Transport mutexes recover poisoning. +If a protocol write or flush fails, the CLI reports the output error and shuts down without +waiting for stdin to close. Output failures are observable only when writing. An active direct +headless export stops at the next photo boundary; a synchronous command and normal backend +cleanup must finish first. Completed exports and library cleanup retain their usual guarantees. + +Library callers serving process stdin use `Server::serve_stdio(stdout())` once per process and +exit after it returns. Its single stdin worker may remain blocked until process exit. The generic +`Server::serve(input, output)` keeps borrowed input scoped: its caller must arrange input EOF or +an input error, even when output fails. + Connect mode and exports inside `command_batch` remain synchronous and do not report MCP progress or cancellation. Use a direct headless export call for this behavior.