//! 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, Write}; use std::sync::atomic::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; type Input = Option>; #[derive(Clone)] pub(super) struct Wire { inbox: Arc>>, out: Sender, deferred: Arc>>, } fn lock(m: &Mutex) -> std::sync::MutexGuard<'_, T> { m.lock().unwrap_or_else(std::sync::PoisonError::into_inner) } impl Server { 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 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 _ = in_tx.send(None); }) { Ok(reader) => reader, Err(error) => { drop(wire); let _ = writer.join(); return Err(error); } }; self.wire = Some(wire.clone()); let served = (|| -> std::io::Result<()> { loop { let queued = lock(&wire.deferred).pop_front(); let line = match queued { Some(line) => line, None => lock(&wire.inbox).recv().ok().flatten(), }; 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"))?; } } })(); self.wire = 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) }) } /// 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 { 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))); closed = true; break; } Ok(None) | Err(TryRecvError::Disconnected) => { lock(&wire.deferred).push_back(None); 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(response(msg["id"].clone(), json!({})).to_string()); } _ => lock(&wire.deferred).push_back(Some(Ok(line))), } } if suppress.load(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(notification.to_string()); } true })) } } #[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)); } }