From 809f7df0fca589d10ad8d02d0e3e7719f14698b4 Mon Sep 17 00:00:00 2001 From: msmps <7691252+msmps@users.noreply.github.com> Date: Sat, 11 Jul 2026 16:31:03 +0100 Subject: [PATCH 1/2] feat(daemon): continuously pump session output --- CONTEXT.md | 11 + crates/pilotty-cli/src/daemon/pty.rs | 101 +-- crates/pilotty-cli/src/daemon/server.rs | 116 ++-- crates/pilotty-cli/src/daemon/session.rs | 770 +++++++++++++++++------ 4 files changed, 694 insertions(+), 304 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index e130709..ac6d041 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -24,6 +24,17 @@ _Avoid_: killed (reserve for explicit `kill`), dead An ended session whose tombstone is gone (TTL, eviction, or daemon restart); indistinguishable from never-existed. +**Finalization**: +The bounded transition from a live session to a tombstone: capture exit status, drain +output to EOF or the deadline, capture final evidence, then stop the session runtime. +_Avoid_: cleanup (too broad), reap (the cleaner's mechanism rather than the transition) + +**Output complete**: +Whether the pump observed PTY EOF before finalization's drain deadline. False means the +final evidence is truthful but may omit later output from a descendant that kept the PTY +open. +_Avoid_: truncated (reserved for retention-ring capacity loss) + ### Observation **Snapshot**: diff --git a/crates/pilotty-cli/src/daemon/pty.rs b/crates/pilotty-cli/src/daemon/pty.rs index 1da0f97..2459e88 100644 --- a/crates/pilotty-cli/src/daemon/pty.rs +++ b/crates/pilotty-cli/src/daemon/pty.rs @@ -4,7 +4,7 @@ use std::io::{Read, Write}; use std::sync::Arc; use anyhow::{Context, Result}; -use portable_pty::{native_pty_system, Child, CommandBuilder, MasterPty, PtySize}; +use portable_pty::{native_pty_system, Child, CommandBuilder, ExitStatus, MasterPty, PtySize}; use tokio::sync::mpsc; use tracing::{debug, error, warn}; @@ -36,7 +36,6 @@ impl From for PtySize { pub struct PtySession { master: Box, child: Box, - size: TermSize, } impl PtySession { @@ -70,7 +69,6 @@ impl PtySession { Ok(Self { master: pair.master, child, - size, }) } @@ -87,11 +85,6 @@ impl PtySession { .take_writer() .context("Failed to take PTY writer") } - /// Get the current terminal size. - pub fn size(&self) -> TermSize { - self.size - } - /// Consume the session and return the master PTY and child process. /// /// Used by AsyncPtyHandle to keep the master for resize operations @@ -111,10 +104,6 @@ const READ_BUFFER_SIZE: usize = 4096; pub struct AsyncPtyHandle { /// Sender for writing to PTY stdin. write_tx: mpsc::Sender>, - /// Receiver for reading from PTY stdout. - /// Wrapped in Mutex for interior mutability so read() can take &self, - /// avoiding the need for &mut self which would require exclusive session access. - read_rx: tokio::sync::Mutex>>, /// Flag to signal shutdown. shutdown: Arc, /// Master PTY for resize operations (sends SIGWINCH). @@ -123,9 +112,6 @@ pub struct AsyncPtyHandle { /// Child process handle for cleanup on shutdown. /// Wrapped in Mutex to allow killing from shutdown(). child: std::sync::Mutex>, - /// Current terminal size, updated on resize. - /// Wrapped in Mutex for interior mutability. - size: std::sync::Mutex, /// Handle to the reader thread for cleanup. reader_thread: Option>, /// Handle to the writer thread for cleanup. @@ -136,10 +122,9 @@ impl AsyncPtyHandle { /// Create async I/O channels for a PTY session. /// /// This spawns background threads for reading and writing to the PTY. - pub fn new(session: PtySession) -> Result { + pub fn new(session: PtySession) -> Result<(Self, mpsc::Receiver>)> { let reader = session.reader()?; let writer = session.writer()?; - let initial_size = session.size(); let (master, child) = session.into_parts(); let shutdown = Arc::new(std::sync::atomic::AtomicBool::new(false)); @@ -159,32 +144,27 @@ impl AsyncPtyHandle { Self::writer_loop(writer, write_rx); }); - Ok(Self { - write_tx, - read_rx: tokio::sync::Mutex::new(read_rx), - shutdown, - master: std::sync::Mutex::new(master), - child: std::sync::Mutex::new(child), - size: std::sync::Mutex::new(initial_size), - reader_thread: Some(reader_thread), - writer_thread: Some(writer_thread), - }) + Ok(( + Self { + write_tx, + shutdown, + master: std::sync::Mutex::new(master), + child: std::sync::Mutex::new(child), + reader_thread: Some(reader_thread), + writer_thread: Some(writer_thread), + }, + read_rx, + )) } /// Resize the PTY and send SIGWINCH to the child process. /// - /// Also updates the internal size tracking. Use `size()` to query. pub fn resize(&self, size: TermSize) -> Result<()> { self.master .lock() .map_err(|_| anyhow::anyhow!("Master PTY mutex poisoned"))? .resize(size.into()) .context("Failed to resize PTY")?; - // Update tracked size - *self - .size - .lock() - .map_err(|_| anyhow::anyhow!("Size mutex poisoned"))? = size; Ok(()) } /// Send bytes to the PTY stdin. @@ -195,30 +175,20 @@ impl AsyncPtyHandle { .context("Failed to send to PTY input channel") } - /// Receive bytes from the PTY stdout. - /// - /// Returns None if the PTY has closed. - pub async fn read(&self) -> Option> { - self.read_rx.lock().await.recv().await - } - - /// Check if the child process has exited without blocking. - /// - /// Returns `Some(true)` if the process has exited, `Some(false)` if still running, - /// or `None` if the mutex is poisoned. - pub fn has_exited(&self) -> Option { + /// Return the child process exit status when it has exited. + pub fn exit_status(&self) -> Result> { self.child .lock() - .ok() - .and_then(|mut child| child.try_wait().ok()) - .map(|status| status.is_some()) + .map_err(|_| anyhow::anyhow!("Child process mutex poisoned"))? + .try_wait() + .context("Failed to inspect child process status") } - /// Shutdown the async PTY I/O and terminate the child process. + /// Signal the I/O threads to stop and terminate the child process. /// - /// This kills the child process to prevent orphaned processes, - /// then signals the I/O threads to stop. - pub async fn shutdown(&self) { + /// The session runtime owns the read receiver and is responsible for ending its + /// pump before dropping this handle. + pub fn terminate(&self) { // Kill the child process first to prevent orphaned processes if let Ok(mut child) = self.child.lock() { if let Err(e) = child.kill() { @@ -237,8 +207,6 @@ impl AsyncPtyHandle { self.shutdown .store(true, std::sync::atomic::Ordering::SeqCst); - // Close the channels to unblock threads - self.read_rx.lock().await.close(); } /// Reader loop running in a background thread. @@ -317,18 +285,16 @@ impl Drop for AsyncPtyHandle { self.shutdown .store(true, std::sync::atomic::Ordering::SeqCst); - // Note: We intentionally don't block on join() here because: + // We intentionally don't block on join() here because: // 1. The reader thread may be blocked on a synchronous read() call // which cannot be interrupted without closing the PTY fd // 2. The threads will terminate on their own when: // - Reader: PTY closes (EOF) or channel is dropped // - Writer: Channel closes when write_tx is dropped // - // The thread handles are stored so they're not detached (which would - // cause issues with thread-local storage), but we don't wait for them. - // This is acceptable because the threads don't hold any resources that - // need explicit cleanup - the channels and PTY handles are cleaned up - // by their own Drop implementations. + // Dropping a std thread handle detaches it. That is acceptable here because the + // threads own no state needed by callers; dropping the PTY and channel handles + // makes them terminate naturally. // Log if threads are still running (helpful for debugging) if let Some(ref handle) = self.reader_thread { @@ -449,14 +415,15 @@ mod tests { let session = PtySession::spawn(&["bash".to_string()], TermSize::default(), None) .expect("Failed to spawn bash"); - let handle = AsyncPtyHandle::new(session).expect("Failed to create async handle"); + let (handle, mut read_rx) = + AsyncPtyHandle::new(session).expect("Failed to create async handle"); // Give bash time to start and print prompt tokio::time::sleep(Duration::from_millis(200)).await; // Drain any initial output (prompt, etc.) while let Ok(Some(_)) = - tokio::time::timeout(Duration::from_millis(100), handle.read()).await + tokio::time::timeout(Duration::from_millis(100), read_rx.recv()).await { // Keep reading until no more output } @@ -470,7 +437,7 @@ mod tests { // The channel should close when bash exits // Try to read, should eventually return None or timeout let _ = tokio::time::timeout(Duration::from_secs(2), async { - while handle.read().await.is_some() { + while read_rx.recv().await.is_some() { // Keep reading until EOF } }) @@ -480,11 +447,7 @@ mod tests { // since we sent exit and bash should have terminated // Shutdown should complete without hanging - let shutdown_result = tokio::time::timeout(Duration::from_secs(2), handle.shutdown()).await; - assert!( - shutdown_result.is_ok(), - "Shutdown timed out, tasks may be stuck" - ); + handle.terminate(); } #[tokio::test] @@ -493,7 +456,7 @@ mod tests { let session = PtySession::spawn(&["sh".to_string()], TermSize { cols: 80, rows: 24 }, None) .expect("spawn"); - let handle = AsyncPtyHandle::new(session).expect("async handle"); + let (handle, _read_rx) = AsyncPtyHandle::new(session).expect("async handle"); // Resize via AsyncPtyHandle (sends SIGWINCH to child process) handle diff --git a/crates/pilotty-cli/src/daemon/server.rs b/crates/pilotty-cli/src/daemon/server.rs index 8ceb88f..b8309ed 100644 --- a/crates/pilotty-cli/src/daemon/server.rs +++ b/crates/pilotty-cli/src/daemon/server.rs @@ -16,7 +16,7 @@ use tokio::task::JoinSet; use tracing::{debug, error, info, warn}; use crate::daemon::paths; -use crate::daemon::session::{SessionId, SessionManager}; +use crate::daemon::session::{ObservationEvent, SessionId, SessionManager}; /// Maximum number of concurrent client connections to prevent resource exhaustion. const MAX_CONNECTIONS: usize = 100; @@ -629,8 +629,8 @@ async fn handle_spawn( } } -/// Poll interval for await_change/settle operations. -const SNAPSHOT_POLL_INTERVAL_MS: u64 = 50; +/// Minimum useful settle window retained for CLI compatibility. +const MIN_SETTLE_MS: u64 = 50; /// Handle snapshot command. /// @@ -658,14 +658,17 @@ async fn handle_snapshot( let format = format.unwrap_or(SnapshotFormat::Full); let with_elements = matches!(format, SnapshotFormat::Full); let timeout = Duration::from_millis(timeout_ms); - // Settle must be at least one poll interval to be meaningful + // Retain the shipped minimum settle window. let settle = Duration::from_millis(if settle_ms > 0 { - settle_ms.max(SNAPSHOT_POLL_INTERVAL_MS) + settle_ms.max(MIN_SETTLE_MS) } else { 0 }); - let poll_interval = Duration::from_millis(SNAPSHOT_POLL_INTERVAL_MS); let start = Instant::now(); + let mut observer = match sessions.observe_session(&session_id).await { + Ok(observer) => observer, + Err(error) => return Response::error(request_id, error), + }; // Phase 1: If await_change is set, wait until content_hash differs if let Some(baseline_hash) = await_change { @@ -684,14 +687,11 @@ async fn handle_snapshot( ); } - let snapshot = match sessions.get_snapshot_data(&session_id, false).await { - Ok(data) => data, - Err(e) => return Response::error(request_id, e), - }; + let snapshot = observer.current(false).await; - if snapshot.content_hash != Some(baseline_hash) { + if snapshot.content_hash != baseline_hash { debug!( - "Screen changed from hash {} to {:?} after {}ms", + "Screen changed from hash {} to {} after {}ms", baseline_hash, snapshot.content_hash, start.elapsed().as_millis() @@ -699,17 +699,24 @@ async fn handle_snapshot( break; } - tokio::time::sleep(poll_interval).await; + let remaining = timeout.saturating_sub(start.elapsed()); + match observer.wait_for_update(remaining).await { + ObservationEvent::Updated => {} + ObservationEvent::OutputClosed => { + tokio::time::sleep(remaining).await; + } + ObservationEvent::Deadline => {} + ObservationEvent::PumpFailed => { + return pump_failure_response(request_id); + } + } } } // Phase 2: If settle_ms > 0, wait for screen stability if settle_ms > 0 { // Get fresh snapshot - don't rely on potentially stale hash from Phase 1 - let snapshot = match sessions.get_snapshot_data(&session_id, false).await { - Ok(data) => data, - Err(e) => return Response::error(request_id, e), - }; + let snapshot = observer.current(false).await; let mut last_hash = snapshot.content_hash; let mut stable_since = Instant::now(); @@ -719,7 +726,7 @@ async fn handle_snapshot( request_id, ApiError::command_failed_with_suggestion( format!( - "Timeout after {}ms waiting for screen to stabilize for {}ms (last hash: {:?})", + "Timeout after {}ms waiting for screen to stabilize for {}ms (last hash: {})", timeout_ms, settle_ms, last_hash @@ -729,15 +736,7 @@ async fn handle_snapshot( ); } - let snapshot = match sessions.get_snapshot_data(&session_id, false).await { - Ok(data) => data, - Err(e) => return Response::error(request_id, e), - }; - - if snapshot.content_hash != last_hash { - last_hash = snapshot.content_hash; - stable_since = Instant::now(); - } else if stable_since.elapsed() >= settle { + if stable_since.elapsed() >= settle { debug!( "Screen stabilized for {}ms after {}ms total", settle_ms, @@ -746,15 +745,32 @@ async fn handle_snapshot( break; } - tokio::time::sleep(poll_interval).await; + let remaining = timeout.saturating_sub(start.elapsed()); + let until_settled = settle.saturating_sub(stable_since.elapsed()); + let wait = remaining.min(until_settled); + match observer.wait_for_update(wait).await { + ObservationEvent::Updated => { + let snapshot = observer.current(false).await; + if snapshot.content_hash != last_hash { + last_hash = snapshot.content_hash; + stable_since = Instant::now(); + } + } + ObservationEvent::OutputClosed => tokio::time::sleep(wait).await, + ObservationEvent::Deadline => {} + ObservationEvent::PumpFailed => { + return pump_failure_response(request_id); + } + } } } // Phase 3: Take final snapshot with requested format - let snapshot = match sessions.get_snapshot_data(&session_id, with_elements).await { - Ok(data) => data, - Err(e) => return Response::error(request_id, e), - }; + let snapshot = observer.current(with_elements).await; + debug!( + "Captured session {} at revision {}", + session_id, snapshot.revision + ); let (cursor_row, cursor_col) = snapshot.cursor_pos; match format { @@ -784,7 +800,7 @@ async fn handle_snapshot( }, text: Some(snapshot.text), elements: snapshot.elements, - content_hash: snapshot.content_hash, + content_hash: Some(snapshot.content_hash), }; Response::success(request_id, ResponseData::ScreenState(screen_state)) } @@ -1183,7 +1199,6 @@ async fn handle_wait_for( ) -> Response { use std::time::{Duration, Instant}; - const POLL_INTERVAL_MS: u64 = 100; let timeout = Duration::from_millis(timeout_ms.unwrap_or(30000)); let use_regex = regex.unwrap_or(false); @@ -1216,8 +1231,11 @@ async fn handle_wait_for( }; let start = Instant::now(); + let mut observer = match sessions.observe_session(&session_id).await { + Ok(observer) => observer, + Err(error) => return Response::error(request_id, error), + }; - // Poll loop loop { // Check timeout first let elapsed = start.elapsed(); @@ -1236,10 +1254,7 @@ async fn handle_wait_for( } // Get current screen text (no elements needed for wait_for) - let snapshot = match sessions.get_snapshot_data(&session_id, false).await { - Ok(data) => data, - Err(e) => return Response::error(request_id, e), - }; + let snapshot = observer.current(false).await; // Check for match let matched = if let Some(ref re) = compiled_regex { @@ -1269,11 +1284,30 @@ async fn handle_wait_for( ); } - // Wait before next poll - tokio::time::sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; + let remaining = timeout.saturating_sub(start.elapsed()); + match observer.wait_for_update(remaining).await { + ObservationEvent::Updated => {} + ObservationEvent::OutputClosed => { + tokio::time::sleep(remaining).await; + } + ObservationEvent::Deadline => {} + ObservationEvent::PumpFailed => { + return pump_failure_response(request_id); + } + } } } +fn pump_failure_response(request_id: &str) -> Response { + Response::error( + request_id, + ApiError::command_failed_with_suggestion( + "Session output pump stopped unexpectedly", + "Inspect daemon logs, then stop and restart the daemon before retrying.", + ), + ) +} + /// Handle shutdown command - gracefully stop the daemon. /// /// Kills all sessions and signals the main run loop to exit. diff --git a/crates/pilotty-cli/src/daemon/session.rs b/crates/pilotty-cli/src/daemon/session.rs index bab876b..6e50ef0 100644 --- a/crates/pilotty-cli/src/daemon/session.rs +++ b/crates/pilotty-cli/src/daemon/session.rs @@ -3,10 +3,12 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use chrono::{DateTime, Utc}; -use tokio::sync::{Mutex, RwLock}; +use portable_pty::ExitStatus; +use tokio::sync::{mpsc, watch, Mutex, RwLock}; +use tokio::task::JoinHandle; use tracing::{debug, info}; use pilotty_core::elements::classify::{detect, ClassifyContext}; @@ -53,41 +55,104 @@ impl From<&str> for SessionId { } } -/// Snapshot data returned from `SessionManager::get_snapshot_data`. -pub struct SnapshotData { - pub text: String, - pub cursor_pos: (u16, u16), - pub cursor_visible: bool, - pub size: TermSize, +/// Atomically observed screen data for a session. +pub(crate) struct SnapshotData { + pub(crate) text: String, + pub(crate) cursor_pos: (u16, u16), + pub(crate) cursor_visible: bool, + pub(crate) size: TermSize, /// Detected UI elements (computed on demand). - pub elements: Option>, + pub(crate) elements: Option>, /// Hash of screen content for change detection. - /// Present when `with_elements=true`. - pub content_hash: Option, + pub(crate) content_hash: u64, + /// Revision matching this exact screen state. + pub(crate) revision: u64, +} + +/// Latest state published by a session's output pump. +#[derive(Debug, Clone, Copy)] +struct PumpState { + revision: u64, + last_output_at: Instant, + output_closed: bool, +} + +/// Terminal state whose content and revision are captured atomically. +struct ObservedTerminal { + emulator: TerminalEmulator, + revision: u64, + size: TermSize, +} + +/// Owned pump task that cannot detach when its session is dropped. +struct PumpTask { + handle: Option>, +} + +impl PumpTask { + fn new(handle: JoinHandle<()>) -> Self { + Self { + handle: Some(handle), + } + } + + async fn abort_and_wait(&mut self) { + let Some(handle) = self.handle.take() else { + return; + }; + handle.abort(); + if let Err(error) = handle.await { + if !error.is_cancelled() { + debug!("Output pump failed during shutdown: {}", error); + } + } + } + + async fn wait(&mut self) { + let Some(handle) = self.handle.take() else { + return; + }; + if let Err(error) = handle.await { + debug!("Output pump failed: {}", error); + } + } +} + +impl Drop for PumpTask { + fn drop(&mut self) { + if let Some(handle) = self.handle.take() { + handle.abort(); + } + } } /// An active PTY session. -pub struct Session { +struct Session { /// Unique session ID. - pub id: SessionId, + id: SessionId, /// Optional human-readable name. - pub name: Option, + name: Option, /// Command that was spawned. - pub command: Vec, + command: Vec, /// When the session was created. - pub created_at: DateTime, - /// Terminal size. - pub size: TermSize, + created_at: DateTime, /// Async handle for PTY I/O. - pub pty: AsyncPtyHandle, - /// Terminal emulator tracking screen state. - /// Wrapped in Mutex for interior mutability (fed by PTY reader task). - pub terminal: Arc>, + pty: AsyncPtyHandle, + observed_terminal: Arc>, + pump_state: watch::Receiver, + pump_task: Mutex, + process_exit: std::sync::Mutex>, +} + +#[derive(Clone)] +struct ProcessExit { + status: ExitStatus, + observed_at: Instant, } impl Session { /// Get session info for protocol responses. - pub fn info(&self) -> SessionInfo { + fn info(&self) -> SessionInfo { SessionInfo { id: self.id.0.clone(), name: self.name.clone(), @@ -97,75 +162,205 @@ impl Session { } /// Check if terminal is in application cursor mode. - pub async fn application_cursor(&self) -> bool { - self.terminal.lock().await.application_cursor() + async fn application_cursor(&self) -> bool { + self.observed_terminal + .lock() + .await + .emulator + .application_cursor() } /// Write bytes to the PTY (send input to the terminal). - pub async fn write(&self, data: &[u8]) -> anyhow::Result<()> { + async fn write(&self, data: &[u8]) -> anyhow::Result<()> { self.pty.write(data).await } - /// Drain pending PTY output and feed to terminal emulator. - /// - /// Call this before taking a snapshot to ensure screen state is current. - /// - /// To prevent blocking on noisy processes, this has limits: - /// - Maximum 100 iterations (reads) - /// - Maximum 1MB total data - /// - 10ms timeout per read - /// - /// These limits ensure the drain completes quickly even with high-output processes. - pub async fn drain_pty_output(&self) { - use std::time::Duration; - use tokio::time::timeout; + async fn snapshot(&self, with_elements: bool) -> SnapshotData { + let terminal = self.observed_terminal.lock().await; + let text = terminal.emulator.get_text(); + let cursor_pos = terminal.emulator.cursor_position(); + let cursor_visible = terminal.emulator.cursor_visible(); + let elements = if with_elements { + let (cursor_row, cursor_col) = cursor_pos; + let context = ClassifyContext::new().with_cursor(cursor_row, cursor_col); + Some(detect(&terminal.emulator, &context)) + } else { + None + }; - // Limits to prevent blocking on noisy processes - const MAX_ITERATIONS: usize = 100; - const MAX_BYTES: usize = 1024 * 1024; // 1 MB + SnapshotData { + content_hash: compute_content_hash(&text), + text, + cursor_pos, + cursor_visible, + size: terminal.size, + elements, + revision: terminal.revision, + } + } - let mut terminal = self.terminal.lock().await; - let mut iterations = 0; - let mut total_bytes = 0; + fn pump_state(&self) -> PumpState { + *self.pump_state.borrow() + } - // Read available data from PTY (non-blocking via short timeout) - loop { - if iterations >= MAX_ITERATIONS { - debug!( - "Drain hit iteration limit ({} iterations, {} bytes)", - iterations, total_bytes - ); - break; + fn observe_process_exit(&self) -> anyhow::Result> { + let mut process_exit = self + .process_exit + .lock() + .map_err(|_| anyhow::anyhow!("Process exit mutex poisoned"))?; + if let Some(exit) = process_exit.as_ref() { + return Ok(Some(exit.clone())); + } + + let Some(status) = self.pty.exit_status()? else { + return Ok(None); + }; + let exit = ProcessExit { + status, + observed_at: Instant::now(), + }; + *process_exit = Some(exit.clone()); + Ok(Some(exit)) + } + + async fn finish_pump(&self, output_complete: bool) { + let mut pump_task = self.pump_task.lock().await; + if output_complete { + pump_task.wait().await; + } else { + pump_task.abort_and_wait().await; + } + } + + async fn wait_for_output_close(&self) -> bool { + let mut pump_state = self.pump_state.clone(); + if pump_state.borrow().output_closed { + return true; + } + + tokio::time::timeout(EXIT_DRAIN_TIMEOUT, async move { + loop { + if pump_state.changed().await.is_err() { + return pump_state.borrow().output_closed; + } + if pump_state.borrow_and_update().output_closed { + return true; + } } + }) + .await + .unwrap_or(false) + } + + async fn shutdown(&self) { + self.pty.terminate(); + let output_complete = self.wait_for_output_close().await; + self.finish_pump(output_complete).await; + } +} - match timeout(Duration::from_millis(10), self.pty.read()).await { - Ok(Some(data)) => { - let len = data.len(); - debug!("Fed {} bytes to terminal emulator", len); - terminal.feed(&data); +/// Result of waiting for a session observation without exposing watch semantics. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ObservationEvent { + Updated, + OutputClosed, + Deadline, + PumpFailed, +} - iterations += 1; - total_bytes += len; +/// Subscription to one session's atomically observed terminal state. +pub(crate) struct SessionObserver { + session: Arc, + pump_state: watch::Receiver, +} - if total_bytes >= MAX_BYTES { - debug!( - "Drain hit byte limit ({} iterations, {} bytes)", - iterations, total_bytes - ); - break; - } +impl SessionObserver { + /// Capture the current screen after marking the current pump state as observed. + /// Output arriving after that mark remains visible to `wait_for_update`. + pub(crate) async fn current(&mut self, with_elements: bool) -> SnapshotData { + { + let _state = self.pump_state.borrow_and_update(); + } + self.session.snapshot(with_elements).await + } + + /// Wait for output, EOF, pump failure, or the supplied duration. + pub(crate) async fn wait_for_update(&mut self, duration: Duration) -> ObservationEvent { + match tokio::time::timeout(duration, self.pump_state.changed()).await { + Err(_) => ObservationEvent::Deadline, + Ok(Ok(())) => { + if self.pump_state.borrow().output_closed { + ObservationEvent::OutputClosed + } else { + ObservationEvent::Updated } - Ok(None) => { - // PTY closed - debug!("PTY channel closed"); - break; + } + Ok(Err(_)) => { + if self.pump_state.borrow().output_closed { + ObservationEvent::OutputClosed + } else { + ObservationEvent::PumpFailed } - Err(_) => { - // Timeout - no more data available + } + } + } +} + +const MAX_PUMP_BATCH_BYTES: usize = 1024 * 1024; +const EXIT_DRAIN_TIMEOUT: Duration = Duration::from_secs(1); + +async fn run_output_pump( + mut read_rx: mpsc::Receiver>, + observed_terminal: Arc>, + state_tx: watch::Sender, +) { + let mut pending = None; + let mut last_output_at = Instant::now(); + + loop { + let first = match pending.take() { + Some(data) => data, + None => match read_rx.recv().await { + Some(data) => data, + None => { + let revision = observed_terminal.lock().await.revision; + state_tx.send_replace(PumpState { + revision, + last_output_at, + output_closed: true, + }); + return; + } + }, + }; + + let mut batch = first; + while batch.len() < MAX_PUMP_BATCH_BYTES { + match read_rx.try_recv() { + Ok(data) if batch.len() + data.len() <= MAX_PUMP_BATCH_BYTES => { + batch.extend_from_slice(&data); + } + Ok(data) => { + pending = Some(data); break; } + Err(mpsc::error::TryRecvError::Empty) => break, + Err(mpsc::error::TryRecvError::Disconnected) => break, } } + + let revision = { + let mut terminal = observed_terminal.lock().await; + terminal.emulator.feed(&batch); + terminal.revision = terminal.revision.saturating_add(1); + terminal.revision + }; + last_output_at = Instant::now(); + state_tx.send_replace(PumpState { + revision, + last_output_at, + output_closed: false, + }); } } @@ -176,7 +371,7 @@ const MAX_SESSIONS: usize = 100; /// /// Thread-safe via interior mutability with RwLock. pub struct SessionManager { - sessions: RwLock>, + sessions: RwLock>>, /// Global snapshot counter for unique snapshot IDs. snapshot_counter: AtomicU64, } @@ -239,36 +434,52 @@ impl SessionManager { // Spawn the PTY session let pty_session = PtySession::spawn(&command, size, cwd.as_deref()) - .map_err(|e| ApiError::spawn_failed(&command, &e.to_string()))?; + .map_err(|error| ApiError::spawn_failed(&command, &format!("{error:#}")))?; - // Wrap in async handle - let pty = AsyncPtyHandle::new(pty_session) - .map_err(|e| ApiError::spawn_failed(&command, &e.to_string()))?; + // Wrap in async handle and transfer sole output ownership to the pump. + let (pty, read_rx) = AsyncPtyHandle::new(pty_session) + .map_err(|error| ApiError::spawn_failed(&command, &format!("{error:#}")))?; - // Create terminal emulator to track screen state - let terminal = Arc::new(Mutex::new(TerminalEmulator::new(size))); + let observed_terminal = Arc::new(Mutex::new(ObservedTerminal { + emulator: TerminalEmulator::new(size), + revision: 0, + size, + })); + let initial_pump_state = PumpState { + revision: 0, + last_output_at: Instant::now(), + output_closed: false, + }; + let (pump_state_tx, pump_state) = watch::channel(initial_pump_state); + let pump_handle = tokio::spawn(run_output_pump( + read_rx, + observed_terminal.clone(), + pump_state_tx, + )); let id = SessionId::new(); - let session = Session { + let session = Arc::new(Session { id: id.clone(), name, command, created_at: Utc::now(), - size, pty, - terminal, - }; + observed_terminal, + pump_state, + pump_task: Mutex::new(PumpTask::new(pump_handle)), + process_exit: std::sync::Mutex::new(None), + }); let mut sessions = self.sessions.write().await; if sessions.len() >= MAX_SESSIONS { drop(sessions); - session.pty.shutdown().await; + session.shutdown().await; return Err(ApiError::session_limit_reached(MAX_SESSIONS)); } if let Some(ref n) = session.name { if sessions.values().any(|s| s.name.as_deref() == Some(n)) { drop(sessions); - session.pty.shutdown().await; + session.shutdown().await; return Err(ApiError::duplicate_session_name(n)); } } @@ -281,7 +492,7 @@ impl SessionManager { /// /// Returns an error if the session doesn't exist. #[cfg(test)] - pub async fn get_session(&self, id: &SessionId, f: F) -> Result + async fn get_session(&self, id: &SessionId, f: F) -> Result where F: FnOnce(&Session) -> R, { @@ -296,15 +507,14 @@ impl SessionManager { /// /// Returns an error if the session doesn't exist. pub async fn kill_session(&self, id: &SessionId) -> Result<(), ApiError> { - let mut sessions = self.sessions.write().await; - match sessions.remove(id) { - Some(session) => { - // Shutdown the async PTY handle (drops writer, signals reader to stop) - session.pty.shutdown().await; - Ok(()) - } - None => Err(ApiError::session_not_found(&id.0)), - } + let session = self + .sessions + .write() + .await + .remove(id) + .ok_or_else(|| ApiError::session_not_found(&id.0))?; + session.shutdown().await; + Ok(()) } /// List all active sessions. @@ -369,67 +579,9 @@ impl SessionManager { } } - /// Get snapshot data for a session. - /// - /// Drains pending PTY output to terminal emulator before capturing snapshot. - /// - /// Uses a read lock on sessions since all operations use interior mutability, - /// avoiding potential deadlocks from holding a write lock during I/O. - /// - /// If `with_elements` is true, element detection runs to identify - /// UI elements like buttons, checkboxes, and menu items. - pub async fn get_snapshot_data( - &self, - id: &SessionId, - with_elements: bool, - ) -> Result { - let sessions = self.sessions.read().await; - let session = sessions - .get(id) - .ok_or_else(|| ApiError::session_not_found(&id.0))?; - - // Drain pending PTY output to update terminal state - session.drain_pty_output().await; - - // Lock terminal once for all reads - let terminal = session.terminal.lock().await; - - // Get snapshot data - let text = terminal.get_text(); - let cursor_pos = terminal.cursor_position(); - let cursor_visible = terminal.cursor_visible(); - let size = session.size; - - // The hash must be computed even when elements are not requested: - // the await_change/settle wait loops poll with `with_elements = false` - // and compare this hash to detect screen changes. - let content_hash = Some(compute_content_hash(&text)); - - // Detect UI elements if requested - let elements = if with_elements { - let (cursor_row, cursor_col) = cursor_pos; - let ctx = ClassifyContext::new().with_cursor(cursor_row, cursor_col); - Some(detect(&*terminal, &ctx)) - } else { - None - }; - - Ok(SnapshotData { - text, - cursor_pos, - cursor_visible, - size, - elements, - content_hash, - }) - } - /// Write bytes to a session's PTY. pub async fn write_to_session(&self, id: &SessionId, data: &[u8]) -> Result<(), ApiError> { - let sessions = self.sessions.read().await; - let session = sessions - .get(id) - .ok_or_else(|| ApiError::session_not_found(&id.0))?; + let session = self.session(id).await?; session .write(data) @@ -446,10 +598,7 @@ impl SessionManager { cols: u16, rows: u16, ) -> Result<(), ApiError> { - let mut sessions = self.sessions.write().await; - let session = sessions - .get_mut(id) - .ok_or_else(|| ApiError::session_not_found(&id.0))?; + let session = self.session(id).await?; let new_size = TermSize { cols, rows }; @@ -459,37 +608,48 @@ impl SessionManager { .resize(new_size) .map_err(|e| ApiError::command_failed(format!("Failed to resize PTY: {}", e)))?; - // Update session's stored size - session.size = new_size; - - // Resize the terminal emulator - session.terminal.lock().await.resize(new_size); + let mut terminal = session.observed_terminal.lock().await; + terminal.size = new_size; + terminal.emulator.resize(new_size); Ok(()) } /// Get terminal size for a session. pub async fn get_terminal_size(&self, id: &SessionId) -> Result { - let sessions = self.sessions.read().await; - let session = sessions - .get(id) - .ok_or_else(|| ApiError::session_not_found(&id.0))?; - Ok(session.size) + let session = self.session(id).await?; + let size = session.observed_terminal.lock().await.size; + Ok(size) } /// Get the application cursor mode for a session. /// - /// Drains pending PTY output first to ensure mode is current. + /// The output pump keeps this mode current. /// When true, arrow keys should send SS3 sequences instead of CSI. pub async fn get_application_cursor_mode(&self, id: &SessionId) -> Result { - let sessions = self.sessions.read().await; - let session = sessions + let session = self.session(id).await?; + Ok(session.application_cursor().await) + } + + async fn session(&self, id: &SessionId) -> Result, ApiError> { + self.sessions + .read() + .await .get(id) - .ok_or_else(|| ApiError::session_not_found(&id.0))?; + .cloned() + .ok_or_else(|| ApiError::session_not_found(&id.0)) + } - // Drain to get current terminal state - session.drain_pty_output().await; - Ok(session.application_cursor().await) + /// Subscribe to screen observations for a live session. + pub(crate) async fn observe_session( + &self, + id: &SessionId, + ) -> Result { + let session = self.session(id).await?; + Ok(SessionObserver { + pump_state: session.pump_state.clone(), + session, + }) } /// Spawn a background task that cleans up dead sessions. @@ -514,33 +674,51 @@ impl SessionManager { break; }; - // Collect IDs of sessions with dead processes - let dead_sessions: Vec = { + // Collect sessions ready to finalize. A direct-process exit starts a + // bounded drain so a descendant cannot keep the session live forever. + let finalizing_sessions: Vec<(SessionId, Arc, bool)> = { let sessions = manager.sessions.read().await; sessions .iter() .filter_map(|(id, session)| { - // Check if child process has exited - if session.pty.has_exited().unwrap_or(false) { - Some(id.clone()) - } else { - None - } + let process_exit = match session.observe_process_exit() { + Ok(exit) => exit, + Err(error) => { + debug!("Failed to inspect session {}: {}", id, error); + return None; + } + }; + let exit = process_exit?; + let pump_state = session.pump_state(); + let output_closed = pump_state.output_closed; + let drain_expired = exit.observed_at.elapsed() >= EXIT_DRAIN_TIMEOUT; + debug!( + "Exit observed for session {} at revision {}, last output {:?} ago", + id, + pump_state.revision, + pump_state.last_output_at.elapsed() + ); + (output_closed || drain_expired) + .then(|| (id.clone(), session.clone(), output_closed)) }) .collect() }; - // Remove dead sessions - if !dead_sessions.is_empty() { + for (id, session, output_complete) in finalizing_sessions { + session.finish_pump(output_complete).await; let mut sessions = manager.sessions.write().await; - for id in dead_sessions { - if let Some(session) = sessions.remove(&id) { - info!( - "Cleaned up session {} ({:?}) - process exited", - id, - session.name.as_deref().unwrap_or("unnamed") - ); - } + let is_same_session = sessions + .get(&id) + .is_some_and(|current| Arc::ptr_eq(current, &session)); + if is_same_session { + sessions.remove(&id); + info!( + "Finalized session {} ({:?}), status: {}, output complete: {}", + id, + session.name.as_deref().unwrap_or("unnamed"), + exit_status_description(&session), + output_complete + ); } } } @@ -548,6 +726,16 @@ impl SessionManager { } } +fn exit_status_description(session: &Session) -> String { + match session.process_exit.lock() { + Ok(exit) => exit + .as_ref() + .map(|exit| exit.status.to_string()) + .unwrap_or_else(|| "unknown".to_string()), + Err(_) => "unavailable".to_string(), + } +} + #[cfg(test)] mod tests { use super::*; @@ -877,6 +1065,200 @@ mod tests { ); } + #[tokio::test] + async fn verbose_session_exits_without_snapshot_reads() { + let manager = Arc::new(SessionManager::new()); + + manager + .create_session( + vec![ + "sh".to_string(), + "-c".to_string(), + "yes output | head -c 1048576".to_string(), + ], + Some("verbose-session".to_string()), + None, + None, + ) + .await + .expect("create verbose session"); + manager.spawn_cleaner(); + + tokio::time::timeout(Duration::from_secs(3), async { + while manager.session_count().await != 0 { + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .expect("verbose session should exit without requiring snapshot reads"); + } + + #[tokio::test] + async fn observer_wakes_when_screen_output_arrives() { + let manager = SessionManager::new(); + let id = manager + .create_session( + vec!["cat".to_string()], + Some("observed-session".to_string()), + None, + None, + ) + .await + .expect("create observed session"); + let mut observer = manager + .observe_session(&id) + .await + .expect("subscribe to session observations"); + let baseline = observer.current(false).await; + + manager + .write_to_session(&id, b"observer-marker\n") + .await + .expect("write marker"); + + assert_eq!( + observer.wait_for_update(Duration::from_secs(1)).await, + ObservationEvent::Updated + ); + let changed = observer.current(false).await; + assert_ne!(changed.content_hash, baseline.content_hash); + assert!(changed.revision > baseline.revision); + assert!(changed.text.contains("observer-marker")); + } + + #[tokio::test] + async fn invisible_output_advances_revision_without_changing_screen_hash() { + let manager = SessionManager::new(); + let id = manager + .create_session( + vec![ + "sh".to_string(), + "-c".to_string(), + r"while :; do printf '\033]0;title\007'; sleep 0.02; done".to_string(), + ], + Some("invisible-output".to_string()), + None, + None, + ) + .await + .expect("create invisible-output session"); + let mut observer = manager + .observe_session(&id) + .await + .expect("subscribe to session observations"); + tokio::time::sleep(Duration::from_millis(100)).await; + let baseline = observer.current(false).await; + + assert_eq!( + observer.wait_for_update(Duration::from_secs(1)).await, + ObservationEvent::Updated + ); + let changed = observer.current(false).await; + assert!(changed.revision > baseline.revision); + assert_eq!(changed.content_hash, baseline.content_hash); + + manager.kill_session(&id).await.expect("kill session"); + } + + #[tokio::test] + async fn output_close_preserves_fast_process_final_screen() { + let manager = SessionManager::new(); + let id = manager + .create_session( + vec![ + "sh".to_string(), + "-c".to_string(), + "printf final-screen-sentinel; exit 3".to_string(), + ], + Some("final-screen".to_string()), + None, + None, + ) + .await + .expect("create fast-exit session"); + let mut observer = manager + .observe_session(&id) + .await + .expect("subscribe to session observations"); + + loop { + match observer.wait_for_update(Duration::from_secs(1)).await { + ObservationEvent::Updated => {} + ObservationEvent::OutputClosed => break, + other => panic!("expected output close, got {other:?}"), + } + } + + let final_screen = observer.current(false).await; + assert!(final_screen.text.contains("final-screen-sentinel")); + manager.kill_session(&id).await.expect("remove session"); + } + + #[tokio::test] + async fn kill_is_not_blocked_by_saturated_input_writers() { + let manager = Arc::new(SessionManager::new()); + let id = manager + .create_session( + vec!["sleep".to_string(), "10".to_string()], + Some("blocked-writers".to_string()), + None, + None, + ) + .await + .expect("create non-reading session"); + + let mut writers = Vec::new(); + for _ in 0..96 { + let manager = manager.clone(); + let id = id.clone(); + writers.push(tokio::spawn(async move { + let data = b"input-line\n".repeat(8 * 1024); + manager.write_to_session(&id, &data).await + })); + } + tokio::time::sleep(Duration::from_millis(100)).await; + + tokio::time::timeout(Duration::from_secs(2), manager.kill_session(&id)) + .await + .expect("kill must not wait for blocked input writers") + .expect("kill session"); + + for writer in writers { + let _write_result = tokio::time::timeout(Duration::from_secs(1), writer) + .await + .expect("input writer should stop after kill") + .expect("input writer task should not panic"); + } + } + + #[tokio::test] + async fn exited_session_finalizes_when_descendant_keeps_pty_open() { + let manager = Arc::new(SessionManager::new()); + + manager + .create_session( + vec![ + "python3".to_string(), + "-c".to_string(), + "import os,signal,time; signal.signal(signal.SIGHUP, signal.SIG_IGN); pid=os.fork(); os._exit(7) if pid else time.sleep(3)".to_string(), + ], + Some("inherited-pty".to_string()), + None, + None, + ) + .await + .expect("create session with descendant"); + manager.spawn_cleaner(); + + tokio::time::timeout(Duration::from_millis(2500), async { + while manager.session_count().await != 0 { + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .expect("an exited session must not stay live while a descendant holds the PTY open"); + } + #[tokio::test] async fn test_is_empty() { let manager = SessionManager::new(); From 0523cb4cf7e783b5a17bf359d739ebaa484715ef Mon Sep 17 00:00:00 2001 From: msmps <7691252+msmps@users.noreply.github.com> Date: Sat, 11 Jul 2026 16:37:59 +0100 Subject: [PATCH 2/2] fix(daemon): cancel pending writes during shutdown Pending PTY input sends retained ended sessions because process termination did not close the bounded input channel consistently across platforms. Signal shutdown directly to every pending writer and reject new writes once termination begins. --- crates/pilotty-cli/src/daemon/pty.rs | 41 ++++++++++++++++++++++++---- 1 file changed, 36 insertions(+), 5 deletions(-) diff --git a/crates/pilotty-cli/src/daemon/pty.rs b/crates/pilotty-cli/src/daemon/pty.rs index 2459e88..b7b9358 100644 --- a/crates/pilotty-cli/src/daemon/pty.rs +++ b/crates/pilotty-cli/src/daemon/pty.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use anyhow::{Context, Result}; use portable_pty::{native_pty_system, Child, CommandBuilder, ExitStatus, MasterPty, PtySize}; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, watch}; use tracing::{debug, error, warn}; /// Terminal size in columns and rows. @@ -104,6 +104,8 @@ const READ_BUFFER_SIZE: usize = 4096; pub struct AsyncPtyHandle { /// Sender for writing to PTY stdin. write_tx: mpsc::Sender>, + /// Wakes pending input sends when session shutdown begins. + write_shutdown: watch::Sender, /// Flag to signal shutdown. shutdown: Arc, /// Master PTY for resize operations (sends SIGWINCH). @@ -131,6 +133,7 @@ impl AsyncPtyHandle { // Create channels let (write_tx, write_rx) = mpsc::channel::>(64); + let (write_shutdown, _write_shutdown_rx) = watch::channel(false); let (read_tx, read_rx) = mpsc::channel::>(64); // Spawn reader thread @@ -147,6 +150,7 @@ impl AsyncPtyHandle { Ok(( Self { write_tx, + write_shutdown, shutdown, master: std::sync::Mutex::new(master), child: std::sync::Mutex::new(child), @@ -169,10 +173,18 @@ impl AsyncPtyHandle { } /// Send bytes to the PTY stdin. pub async fn write(&self, data: &[u8]) -> Result<()> { - self.write_tx - .send(data.to_vec()) - .await - .context("Failed to send to PTY input channel") + let mut shutdown = self.write_shutdown.subscribe(); + if *shutdown.borrow() { + anyhow::bail!("PTY input is shutting down"); + } + + tokio::select! { + biased; + _ = shutdown.changed() => anyhow::bail!("PTY input is shutting down"), + result = self.write_tx.send(data.to_vec()) => { + result.context("Failed to send to PTY input channel") + } + } } /// Return the child process exit status when it has exited. @@ -189,6 +201,8 @@ impl AsyncPtyHandle { /// The session runtime owns the read receiver and is responsible for ending its /// pump before dropping this handle. pub fn terminate(&self) { + self.write_shutdown.send_replace(true); + // Kill the child process first to prevent orphaned processes if let Ok(mut child) = self.child.lock() { if let Err(e) = child.kill() { @@ -265,6 +279,8 @@ impl AsyncPtyHandle { impl Drop for AsyncPtyHandle { fn drop(&mut self) { + self.write_shutdown.send_replace(true); + // Kill the child process first to prevent orphaned processes. // This mirrors the logic in shutdown() but is synchronous since Drop can't be async. if let Ok(mut child) = self.child.lock() { @@ -450,6 +466,21 @@ mod tests { handle.terminate(); } + #[tokio::test] + async fn write_is_rejected_after_terminate() { + let session = + PtySession::spawn(&["cat".to_string()], TermSize::default(), None).expect("spawn cat"); + let (handle, _read_rx) = AsyncPtyHandle::new(session).expect("create async handle"); + + handle.terminate(); + + let error = handle + .write(b"must-not-be-queued") + .await + .expect_err("writes must stop when PTY shutdown begins"); + assert!(error.to_string().contains("shutting down")); + } + #[tokio::test] async fn test_async_pty_handle_resize() { // Spawn a shell