//! Stdio progress and cancellation between photos. The backend keeps its session on the //! serving thread; only stdin and stdout move to transport threads. Other calls wait in order. use std::collections::VecDeque; 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}; use serde_json::{Value, json}; use super::{Server, response}; use crate::backend::ProgressHook; enum Input { Line(std::io::Result), Closed, OutputFailed, } #[derive(Clone)] pub(super) struct Wire { inbox: Arc>>, out: Sender>, failed: Arc, deferred: Arc>>, } 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 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; 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); } 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 => { if wire.failed.load(Ordering::Acquire) { return Ok(()); } lock(&wire.inbox).recv().map_err(|_| std::io::Error::other("MCP input disconnected without EOF"))? } }; 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); written.and(served) }) } /// Called at photo boundaries, including completion. EOF finishes the pending export; /// cancellation stops before the next photo. No file scan or output deletion is needed: /// the existing writer replaces each whole file atomically. pub(super) fn progress_hook(&self) -> Option { let wire = self.wire.clone()?; let (id, token) = self.current.clone()?; let suppress = self.suppress.clone(); let mut closed = false; let mut last = None; 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(Input::Line(Ok(line))) => line, Ok(Input::Line(Err(error))) => { lock(&wire.deferred).push_back(Input::Line(Err(error))); closed = true; break; } 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; } Err(TryRecvError::Empty) => break, }; let msg: Value = serde_json::from_str(&line).unwrap_or(Value::Null); match msg.get("method").and_then(Value::as_str) { Some("notifications/cancelled") if msg.get("id").is_none() && msg["params"]["requestId"] == id => { suppress.store(true, Ordering::Relaxed); } Some("ping") 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(Some(response(msg["id"].clone(), json!({})).to_string())); } _ => lock(&wire.deferred).push_back(Input::Line(Ok(line))), } } if suppress.load(Ordering::Relaxed) || wire.failed.load(Ordering::Acquire) { suppress.store(true, Ordering::Relaxed); return false; } if let Some(token) = &token && total > 0 && last.is_none_or(|last| done > last) && (done == total || reported.is_none_or(|at| at.elapsed() >= Duration::from_millis(100))) { last = Some(done); reported = Some(Instant::now()); let message = if done < total { format!("Exporting {name} ({} of {total})", done.saturating_add(1)) } else { format!("Exported {total} photo(s)") }; let notification = json!({"jsonrpc":"2.0","method":"notifications/progress", "params":{"progressToken":token,"progress":done,"total":total,"message":message}}); let _ = wire.out.send(Some(notification.to_string())); } true })) } } #[cfg(test)] #[path = "progress_tests.rs"] mod shutdown_tests; #[cfg(test)] mod tests { use super::*; #[test] fn transport_lock_recovers_after_panic() { let state = Arc::new(Mutex::new(VecDeque::new())); let held = state.clone(); let _ = std::thread::spawn(move || { held.lock().unwrap().push_back(7); let _guard = held.lock().unwrap(); panic!("synthetic transport panic"); }) .join(); assert_eq!(lock(&state).pop_front(), Some(7)); lock(&state).push_back(8); assert_eq!(lock(&state).pop_front(), Some(8)); } }