diff --git a/README.md b/README.md index 307e9f8..5bbecea 100644 --- a/README.md +++ b/README.md @@ -107,6 +107,7 @@ pilotty stop pilotty spawn # Spawn a TUI app (e.g., pilotty spawn vim file.txt) pilotty spawn --name myapp # Spawn with a custom session name pilotty spawn --cwd /path cmd # Spawn in a specific working directory +pilotty spawn --retain-bytes 1048576 # Override retained output limit pilotty kill # Kill default session pilotty kill -s myapp # Kill specific session pilotty list-sessions # List all active sessions @@ -115,6 +116,18 @@ pilotty daemon # Manually start daemon (usually auto-starts) pilotty examples # Show end-to-end workflow example ``` +### Retained Output + +```bash +pilotty logs # Raw output bytes from the default session +pilotty logs -s myapp # Raw output bytes from a named session +pilotty logs -s myapp > output.bin # Save raw ANSI evidence without metadata +``` + +`logs` writes retained output bytes to stdout and exact retention accounting to stderr. +Retention is bounded to 2 MiB per session by default, and reports when older bytes were +dropped. + ### Screen Capture ```bash @@ -397,6 +410,7 @@ All errors include AI-friendly suggestions: |----------|-------------| | `PILOTTY_SESSION` | Default session name | | `PILOTTY_SOCKET_DIR` | Override socket directory | +| `PILOTTY_RETAIN_BYTES` | Default retained raw output bytes per session (default: 2 MiB) | | `RUST_LOG` | Logging level (e.g., `debug`, `info`) | ## Usage with AI Agents diff --git a/crates/pilotty-cli/src/args.rs b/crates/pilotty-cli/src/args.rs index 26ff834..0549067 100644 --- a/crates/pilotty-cli/src/args.rs +++ b/crates/pilotty-cli/src/args.rs @@ -95,6 +95,9 @@ Examples: /// List all active sessions ListSessions, + /// Print retained raw output for a session + Logs(LogsArgs), + /// Resize the terminal Resize(ResizeArgs), @@ -135,6 +138,10 @@ pub struct SpawnArgs { /// Working directory for the spawned process [default: current directory] #[arg(long, value_name = "DIR")] pub cwd: Option, + + /// Maximum raw output bytes retained for this session + #[arg(long, value_name = "BYTES")] + pub retain_bytes: Option, } #[derive(Debug, clap::Args)] @@ -144,6 +151,13 @@ pub struct KillArgs { pub session: Option, } +#[derive(Debug, clap::Args)] +pub struct LogsArgs { + /// Target session by name or ID [default: default] + #[arg(short, long, help = SESSION_HELP)] + pub session: Option, +} + #[derive(Debug, clap::Args)] pub struct SnapshotArgs { /// Output format @@ -305,4 +319,24 @@ mod tests { _ => panic!("Expected spawn command"), } } + + #[test] + fn spawn_parses_retention_override() { + let cli = Cli::parse_from(["pilotty", "spawn", "--retain-bytes", "4096", "sh"]); + + match cli.command { + Commands::Spawn(args) => assert_eq!(args.retain_bytes, Some(4096)), + _ => panic!("Expected spawn command"), + } + } + + #[test] + fn logs_parses_session_target() { + let cli = Cli::parse_from(["pilotty", "logs", "--session", "editor"]); + + match cli.command { + Commands::Logs(args) => assert_eq!(args.session.as_deref(), Some("editor")), + _ => panic!("Expected logs command"), + } + } } diff --git a/crates/pilotty-cli/src/daemon/client.rs b/crates/pilotty-cli/src/daemon/client.rs index a364dbb..7c71c20 100644 --- a/crates/pilotty-cli/src/daemon/client.rs +++ b/crates/pilotty-cli/src/daemon/client.rs @@ -235,7 +235,8 @@ impl DaemonClient { mod tests { use super::*; use crate::daemon::server::DaemonServer; - use pilotty_core::protocol::{Command, PROTOCOL_VERSION}; + use pilotty_core::protocol::{Command, ResponseData, PROTOCOL_VERSION}; + use tokio::net::UnixListener; #[tokio::test] async fn test_client_connects_to_running_daemon() { @@ -304,4 +305,66 @@ mod tests { ProtocolAction::Send ); } + + #[tokio::test] + async fn old_daemon_is_probed_and_logs_is_rejected_before_transmission() { + let socket_path = + std::env::temp_dir().join(format!("pilotty-probe-{}.sock", std::process::id())); + let _ = std::fs::remove_file(&socket_path); + let listener = UnixListener::bind(&socket_path).expect("bind protocol fixture"); + let peer = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept client"); + let (reader, mut writer) = stream.into_split(); + let mut reader = BufReader::new(reader); + let mut line = String::new(); + reader.read_line(&mut line).await.expect("read probe"); + let probe: Request = serde_json::from_str(&line).expect("parse probe"); + assert_eq!(probe.command, Command::ListSessions); + + let response = Response { + id: probe.id, + success: true, + data: Some(ResponseData::Sessions { sessions: vec![] }), + error: None, + protocol: LEGACY_PROTOCOL_VERSION, + }; + writer + .write_all( + serde_json::to_string(&response) + .expect("serialize response") + .as_bytes(), + ) + .await + .expect("write response"); + writer.write_all(b"\n").await.expect("write newline"); + writer.flush().await.expect("flush response"); + + line.clear(); + assert!( + timeout(Duration::from_millis(100), reader.read_line(&mut line)) + .await + .is_err(), + "version-gated command must not reach an old daemon" + ); + }); + + let stream = UnixStream::connect(&socket_path) + .await + .expect("connect to protocol fixture"); + let mut client = DaemonClient { + stream, + daemon_protocol: None, + }; + let error = client + .request(Request::new( + "logs-request", + Command::Logs { session: None }, + )) + .await + .expect_err("old daemon must be rejected"); + + assert!(error.to_string().contains("protocol 1")); + peer.await.expect("protocol fixture task"); + let _ = std::fs::remove_file(&socket_path); + } } diff --git a/crates/pilotty-cli/src/daemon/mod.rs b/crates/pilotty-cli/src/daemon/mod.rs index c4098a1..ecfa133 100644 --- a/crates/pilotty-cli/src/daemon/mod.rs +++ b/crates/pilotty-cli/src/daemon/mod.rs @@ -3,6 +3,7 @@ pub mod client; pub mod paths; pub mod pty; +pub mod retention; pub mod server; pub mod session; pub mod terminal; diff --git a/crates/pilotty-cli/src/daemon/retention.rs b/crates/pilotty-cli/src/daemon/retention.rs new file mode 100644 index 0000000..54de23a --- /dev/null +++ b/crates/pilotty-cli/src/daemon/retention.rs @@ -0,0 +1,120 @@ +//! Bounded retention of raw PTY output. + +use std::collections::VecDeque; + +/// Default number of raw output bytes retained for each session. +pub(crate) const DEFAULT_RETAIN_BYTES: usize = 2 * 1024 * 1024; + +/// A point-in-time copy of a session's retained raw output and accounting. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct RetentionSnapshot { + pub(crate) bytes: Vec, + pub(crate) total_bytes: u64, + pub(crate) retained_bytes: u64, + pub(crate) dropped_bytes: u64, + pub(crate) truncated: bool, +} + +/// Bounded tail of raw PTY output. +pub(crate) struct RetentionRing { + bytes: VecDeque, + capacity: usize, + total_bytes: u64, +} + +impl RetentionRing { + pub(crate) fn new(capacity: usize) -> Self { + Self { + bytes: VecDeque::new(), + capacity, + total_bytes: 0, + } + } + + pub(crate) fn append(&mut self, output: &[u8]) { + let output_len = u64::try_from(output.len()).unwrap_or(u64::MAX); + self.total_bytes = self.total_bytes.saturating_add(output_len); + + if self.capacity == 0 { + self.bytes.clear(); + return; + } + + if output.len() >= self.capacity { + self.bytes.clear(); + self.bytes + .extend(output[output.len() - self.capacity..].iter().copied()); + return; + } + + let excess = self + .bytes + .len() + .saturating_add(output.len()) + .saturating_sub(self.capacity); + if excess > 0 { + self.bytes.drain(..excess); + } + self.bytes.extend(output.iter().copied()); + } + + pub(crate) fn snapshot(&self) -> RetentionSnapshot { + let bytes: Vec = self.bytes.iter().copied().collect(); + let retained_bytes = u64::try_from(bytes.len()).unwrap_or(u64::MAX); + let dropped_bytes = self.total_bytes.saturating_sub(retained_bytes); + RetentionSnapshot { + bytes, + total_bytes: self.total_bytes, + retained_bytes, + dropped_bytes, + truncated: dropped_bytes > 0, + } + } +} + +#[cfg(test)] +mod tests { + use crate::daemon::retention::RetentionRing; + + #[test] + fn retains_only_the_newest_bytes_with_exact_accounting() { + let mut retention = RetentionRing::new(5); + + retention.append(b"abc"); + retention.append(b"defg"); + + let snapshot = retention.snapshot(); + assert_eq!(snapshot.bytes, b"cdefg"); + assert_eq!(snapshot.total_bytes, 7); + assert_eq!(snapshot.retained_bytes, 5); + assert_eq!(snapshot.dropped_bytes, 2); + assert!(snapshot.truncated); + } + + #[test] + fn oversized_append_keeps_only_its_tail() { + let mut retention = RetentionRing::new(4); + + retention.append(b"old"); + retention.append(b"123456"); + + let snapshot = retention.snapshot(); + assert_eq!(snapshot.bytes, b"3456"); + assert_eq!(snapshot.total_bytes, 9); + assert_eq!(snapshot.dropped_bytes, 5); + } + + #[test] + fn zero_capacity_counts_every_byte_as_dropped() { + let mut retention = RetentionRing::new(0); + + retention.append(b"evidence"); + + let snapshot = retention.snapshot(); + assert!(snapshot.bytes.is_empty()); + assert_eq!(snapshot.total_bytes, 8); + assert_eq!(snapshot.retained_bytes, 0); + assert_eq!(snapshot.dropped_bytes, 8); + assert!(snapshot.truncated); + } +} diff --git a/crates/pilotty-cli/src/daemon/server.rs b/crates/pilotty-cli/src/daemon/server.rs index 3e4b352..6fd723b 100644 --- a/crates/pilotty-cli/src/daemon/server.rs +++ b/crates/pilotty-cli/src/daemon/server.rs @@ -18,8 +18,11 @@ use tokio::task::JoinSet; use tracing::{debug, error, info, warn}; use crate::daemon::paths; +use crate::daemon::retention::DEFAULT_RETAIN_BYTES; use crate::daemon::session::{ObservationEvent, SessionId, SessionManager}; +const RETAIN_BYTES_ENV: &str = "PILOTTY_RETAIN_BYTES"; + /// Maximum number of concurrent client connections to prevent resource exhaustion. const MAX_CONNECTIONS: usize = 100; @@ -53,7 +56,8 @@ impl DaemonServer { paths::ensure_socket_dir().context("Failed to create socket directory")?; let socket_path = paths::get_socket_path(None); let pid_path = paths::get_pid_path(None); - Self::bind_to(socket_path, pid_path).await + let retain_bytes = retain_bytes_from_env()?; + Self::bind_to_with_retain_bytes(socket_path, pid_path, retain_bytes).await } /// Create a new daemon server bound to a specific socket path. @@ -63,7 +67,16 @@ impl DaemonServer { /// 2. If socket in use, check PID file to see if daemon is alive /// 3. If daemon dead, remove stale socket and retry /// 4. If daemon alive, return error + #[cfg(test)] pub async fn bind_to(socket_path: PathBuf, pid_path: PathBuf) -> Result { + Self::bind_to_with_retain_bytes(socket_path, pid_path, DEFAULT_RETAIN_BYTES).await + } + + pub(crate) async fn bind_to_with_retain_bytes( + socket_path: PathBuf, + pid_path: PathBuf, + retain_bytes: usize, + ) -> Result { if let Some(parent) = socket_path.parent() { std::fs::create_dir_all(parent).with_context(|| { format!("Failed to create socket directory for {:?}", socket_path) @@ -142,7 +155,7 @@ impl DaemonServer { listener, socket_path, pid_path, - sessions: Arc::new(SessionManager::new()), + sessions: Arc::new(SessionManager::with_default_retain_bytes(retain_bytes)), connection_semaphore: Arc::new(Semaphore::new(MAX_CONNECTIONS)), shutdown: Arc::new(Notify::new()), }) @@ -316,6 +329,22 @@ impl DaemonServer { } } +fn retain_bytes_from_env() -> Result { + parse_retain_bytes(std::env::var_os(RETAIN_BYTES_ENV).as_deref()) +} + +fn parse_retain_bytes(value: Option<&std::ffi::OsStr>) -> Result { + let Some(value) = value else { + return Ok(DEFAULT_RETAIN_BYTES); + }; + let value = value + .to_str() + .with_context(|| format!("{RETAIN_BYTES_ENV} must contain a non-negative integer"))?; + value.parse::().with_context(|| { + format!("{RETAIN_BYTES_ENV} must be a non-negative integer number of bytes, got '{value}'") + }) +} + /// Kill all active sessions during shutdown. /// /// Used by both the shutdown command handler and the idle shutdown task. @@ -521,7 +550,18 @@ async fn handle_request( command, session_name, cwd, - } => handle_spawn(&request_id, &sessions, command, session_name, cwd).await, + retain_bytes, + } => { + handle_spawn( + &request_id, + &sessions, + command, + session_name, + cwd, + retain_bytes, + ) + .await + } Command::Snapshot { session, @@ -544,6 +584,8 @@ async fn handle_request( Command::ListSessions => handle_list_sessions(&request_id, &sessions).await, + Command::Logs { session } => handle_logs(&request_id, &sessions, session).await, + Command::Kill { session } => handle_kill(&request_id, &sessions, session).await, Command::Type { text, session } => handle_type(&request_id, &sessions, text, session).await, @@ -602,6 +644,7 @@ async fn handle_spawn( command: Vec, session_name: Option, cwd: Option, + retain_bytes: Option, ) -> Response { if command.is_empty() { return Response::error( @@ -636,8 +679,21 @@ async fn handle_spawn( } } + let retain_bytes = match retain_bytes.map(usize::try_from).transpose() { + Ok(retain_bytes) => retain_bytes, + Err(_) => { + return Response::error( + request_id, + ApiError::invalid_input_with_suggestion( + "Retention limit is too large for this daemon", + "Choose a smaller --retain-bytes value.", + ), + ); + } + }; + match sessions - .create_session(command.clone(), session_name, None, cwd) + .create_session_with_retention(command.clone(), session_name, None, cwd, retain_bytes) .await { Ok(id) => { @@ -654,6 +710,32 @@ async fn handle_spawn( } } +/// Handle logs command. +async fn handle_logs( + request_id: &str, + sessions: &SessionManager, + session: Option, +) -> Response { + let session_id = match sessions.resolve_session(session.as_deref()).await { + Ok(id) => id, + Err(error) => return Response::error(request_id, error), + }; + + match sessions.session_logs(&session_id).await { + Ok(logs) => Response::success( + request_id, + ResponseData::Logs { + bytes: logs.bytes, + total_bytes: logs.total_bytes, + retained_bytes: logs.retained_bytes, + dropped_bytes: logs.dropped_bytes, + truncated: logs.truncated, + }, + ), + Err(error) => Response::error(request_id, error), + } +} + /// Minimum useful settle window retained for CLI compatibility. const MIN_SETTLE_MS: u64 = 50; @@ -1375,6 +1457,129 @@ mod tests { use tokio::time::timeout; use uuid::Uuid; + #[test] + fn retention_environment_value_is_validated() { + assert_eq!(parse_retain_bytes(None).expect("default"), 2 * 1024 * 1024); + assert_eq!( + parse_retain_bytes(Some(std::ffi::OsStr::new("4096"))).expect("configured value"), + 4096 + ); + assert!(parse_retain_bytes(Some(std::ffi::OsStr::new("many"))).is_err()); + } + + #[tokio::test] + async fn legacy_request_cannot_dispatch_logs() { + let response = handle_request( + Request { + id: "legacy-logs".to_string(), + command: Command::Logs { session: None }, + protocol: 0, + }, + Arc::new(SessionManager::new()), + Arc::new(Notify::new()), + ) + .await; + + assert!(!response.success); + assert_eq!(response.protocol, PROTOCOL_VERSION); + assert!(matches!( + response.error.map(|error| error.code), + Some(ErrorCode::InvalidInput) + )); + } + + #[tokio::test] + async fn logs_returns_bounded_ordered_raw_output_over_the_socket() { + let temp_dir = std::env::temp_dir(); + let socket_path = temp_dir.join(format!("pilotty-logs-{}.sock", std::process::id())); + let pid_path = socket_path.with_extension("pid"); + let server = DaemonServer::bind_to_with_retain_bytes(socket_path.clone(), pid_path, 4) + .await + .expect("bind server"); + let server_handle = tokio::spawn(async move { + let _ = timeout(Duration::from_secs(3), server.run()).await; + }); + + let stream = UnixStream::connect(&socket_path) + .await + .expect("connect to server"); + let (reader, mut writer) = stream.into_split(); + let mut reader = BufReader::new(reader); + let spawn = Request::new( + "spawn-logs", + Command::Spawn { + command: vec![ + "sh".to_string(), + "-c".to_string(), + "printf abcdef; sleep 2".to_string(), + ], + session_name: Some("logs-test".to_string()), + cwd: None, + retain_bytes: None, + }, + ); + writer + .write_all( + serde_json::to_string(&spawn) + .expect("serialize spawn") + .as_bytes(), + ) + .await + .expect("write spawn"); + writer.write_all(b"\n").await.expect("write newline"); + writer.flush().await.expect("flush spawn"); + let mut line = String::new(); + reader.read_line(&mut line).await.expect("read spawn"); + let spawn_response: Response = serde_json::from_str(&line).expect("parse spawn response"); + assert!(spawn_response.success); + + let logs = timeout(Duration::from_secs(2), async { + loop { + let request = Request::new( + Uuid::new_v4().to_string(), + Command::Logs { + session: Some("logs-test".to_string()), + }, + ); + writer + .write_all( + serde_json::to_string(&request) + .expect("serialize logs") + .as_bytes(), + ) + .await + .expect("write logs"); + writer.write_all(b"\n").await.expect("write newline"); + writer.flush().await.expect("flush logs"); + line.clear(); + reader.read_line(&mut line).await.expect("read logs"); + let response: Response = serde_json::from_str(&line).expect("parse logs response"); + if matches!( + &response.data, + Some(ResponseData::Logs { total_bytes: 6, .. }) + ) { + break response; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("retained output"); + + assert_eq!( + logs.data, + Some(ResponseData::Logs { + bytes: b"cdef".to_vec(), + total_bytes: 6, + retained_bytes: 4, + dropped_bytes: 2, + truncated: true, + }) + ); + server_handle.abort(); + let _ = std::fs::remove_file(&socket_path); + } + #[tokio::test] async fn test_daemon_accepts_and_responds() { // Use a temp socket path @@ -1619,6 +1824,7 @@ mod tests { command: vec!["echo".to_string(), "hello from test".to_string()], session_name: Some("test-snap".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -1727,6 +1933,7 @@ mod tests { command: vec!["echo".to_string(), "full format test".to_string()], session_name: Some("full-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -1846,6 +2053,7 @@ mod tests { command: vec!["cat".to_string()], session_name: Some("type-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -1967,6 +2175,7 @@ mod tests { command: vec!["cat".to_string()], session_name: Some("key-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -2130,6 +2339,7 @@ mod tests { command: vec!["cat".to_string()], session_name: Some("click-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -2227,6 +2437,7 @@ mod tests { command: vec!["cat".to_string()], session_name: Some("scroll-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -2335,6 +2546,7 @@ mod tests { command: vec!["echo".to_string(), "hello world marker".to_string()], session_name: Some("waitfor-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -2428,6 +2640,7 @@ mod tests { command: vec!["echo".to_string(), "version 1.2.3 ready".to_string()], session_name: Some("waitfor-re-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -2520,6 +2733,7 @@ mod tests { command: vec!["cat".to_string()], session_name: Some("waitfor-to-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -2610,6 +2824,7 @@ mod tests { command: vec!["echo".to_string(), "test".to_string()], session_name: Some("waitfor-bad-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -2701,6 +2916,7 @@ mod tests { command: vec!["cat".to_string()], session_name: Some("await-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -2869,6 +3085,7 @@ mod tests { command: vec!["cat".to_string()], session_name: Some("await-static-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -3005,6 +3222,7 @@ mod tests { ], session_name: Some("settle-busy-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -3103,6 +3321,7 @@ mod tests { command: vec!["pwd".to_string()], session_name: Some("bad-cwd-test".to_string()), cwd: Some("/nonexistent/path/that/does/not/exist".to_string()), + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); @@ -3180,6 +3399,7 @@ mod tests { ], session_name: Some("elem-test".to_string()), cwd: None, + retain_bytes: None, }, }; let request_json = serde_json::to_string(&spawn_request).unwrap(); diff --git a/crates/pilotty-cli/src/daemon/session.rs b/crates/pilotty-cli/src/daemon/session.rs index 6e50ef0..41b31c4 100644 --- a/crates/pilotty-cli/src/daemon/session.rs +++ b/crates/pilotty-cli/src/daemon/session.rs @@ -18,6 +18,7 @@ use pilotty_core::protocol::SessionInfo; use pilotty_core::snapshot::compute_content_hash; use crate::daemon::pty::{AsyncPtyHandle, PtySession, TermSize}; +use crate::daemon::retention::{RetentionRing, RetentionSnapshot, DEFAULT_RETAIN_BYTES}; use crate::daemon::terminal::TerminalEmulator; /// Unique identifier for a session. @@ -138,6 +139,7 @@ struct Session { created_at: DateTime, /// Async handle for PTY I/O. pty: AsyncPtyHandle, + retention: Arc>, observed_terminal: Arc>, pump_state: watch::Receiver, pump_task: Mutex, @@ -311,6 +313,7 @@ const EXIT_DRAIN_TIMEOUT: Duration = Duration::from_secs(1); async fn run_output_pump( mut read_rx: mpsc::Receiver>, + retention: Arc>, observed_terminal: Arc>, state_tx: watch::Sender, ) { @@ -349,6 +352,7 @@ async fn run_output_pump( } } + retention.lock().await.append(&batch); let revision = { let mut terminal = observed_terminal.lock().await; terminal.emulator.feed(&batch); @@ -372,6 +376,7 @@ const MAX_SESSIONS: usize = 100; /// Thread-safe via interior mutability with RwLock. pub struct SessionManager { sessions: RwLock>>, + default_retain_bytes: usize, /// Global snapshot counter for unique snapshot IDs. snapshot_counter: AtomicU64, } @@ -385,8 +390,14 @@ impl Default for SessionManager { impl SessionManager { /// Create a new session manager. pub fn new() -> Self { + Self::with_default_retain_bytes(DEFAULT_RETAIN_BYTES) + } + + /// Create a session manager with an explicit default retention limit. + pub(crate) fn with_default_retain_bytes(default_retain_bytes: usize) -> Self { Self { sessions: RwLock::new(HashMap::new()), + default_retain_bytes, snapshot_counter: AtomicU64::new(1), } } @@ -407,12 +418,25 @@ impl SessionManager { /// - Spawn fails /// /// If `cwd` is provided, the spawned process runs in that directory. + #[cfg(test)] pub async fn create_session( &self, command: Vec, name: Option, size: Option, cwd: Option, + ) -> Result { + self.create_session_with_retention(command, name, size, cwd, None) + .await + } + + pub async fn create_session_with_retention( + &self, + command: Vec, + name: Option, + size: Option, + cwd: Option, + retain_bytes: Option, ) -> Result { let name = name.or_else(|| Some("default".to_string())); @@ -440,6 +464,9 @@ impl SessionManager { let (pty, read_rx) = AsyncPtyHandle::new(pty_session) .map_err(|error| ApiError::spawn_failed(&command, &format!("{error:#}")))?; + let retention = Arc::new(Mutex::new(RetentionRing::new( + retain_bytes.unwrap_or(self.default_retain_bytes), + ))); let observed_terminal = Arc::new(Mutex::new(ObservedTerminal { emulator: TerminalEmulator::new(size), revision: 0, @@ -453,6 +480,7 @@ impl SessionManager { let (pump_state_tx, pump_state) = watch::channel(initial_pump_state); let pump_handle = tokio::spawn(run_output_pump( read_rx, + retention.clone(), observed_terminal.clone(), pump_state_tx, )); @@ -464,6 +492,7 @@ impl SessionManager { command, created_at: Utc::now(), pty, + retention, observed_terminal, pump_state, pump_task: Mutex::new(PumpTask::new(pump_handle)), @@ -589,6 +618,13 @@ impl SessionManager { .map_err(|e| ApiError::write_failed(&e.to_string())) } + /// Capture the retained raw output and its exact accounting. + pub(crate) async fn session_logs(&self, id: &SessionId) -> Result { + let session = self.session(id).await?; + let snapshot = session.retention.lock().await.snapshot(); + Ok(snapshot) + } + /// Resize a session's terminal. /// /// Updates both the PTY size (sends SIGWINCH to child) and the terminal emulator. @@ -1259,6 +1295,75 @@ mod tests { .expect("an exited session must not stay live while a descendant holds the PTY open"); } + #[tokio::test] + async fn configured_default_is_injected_into_new_sessions() { + let manager = SessionManager::with_default_retain_bytes(4); + let id = manager + .create_session_with_retention( + vec!["printf".to_string(), "abcdef".to_string()], + Some("default-retention".to_string()), + None, + None, + None, + ) + .await + .expect("create session"); + + let logs = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let logs = manager.session_logs(&id).await.expect("read logs"); + if logs.total_bytes == 6 { + break logs; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("session output"); + + assert_eq!(logs.bytes, b"cdef"); + assert_eq!(logs.retained_bytes, 4); + assert_eq!(logs.dropped_bytes, 2); + assert!(logs.truncated); + manager.kill_session(&id).await.expect("remove session"); + } + + #[tokio::test] + async fn session_override_retains_ordered_raw_output() { + let manager = SessionManager::with_default_retain_bytes(2); + let expected = b"\x1b[31mred\x1b[0m"; + let id = manager + .create_session_with_retention( + vec![ + "printf".to_string(), + String::from_utf8(expected.to_vec()).expect("valid test bytes"), + ], + Some("retention-override".to_string()), + None, + None, + Some(expected.len()), + ) + .await + .expect("create session"); + + let logs = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let logs = manager.session_logs(&id).await.expect("read logs"); + if logs.total_bytes == expected.len() as u64 { + break logs; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("session output"); + + assert_eq!(logs.bytes, expected); + assert_eq!(logs.dropped_bytes, 0); + assert!(!logs.truncated); + manager.kill_session(&id).await.expect("remove session"); + } + #[tokio::test] async fn test_is_empty() { let manager = SessionManager::new(); diff --git a/crates/pilotty-cli/src/main.rs b/crates/pilotty-cli/src/main.rs index 2495fa3..598e58a 100644 --- a/crates/pilotty-cli/src/main.rs +++ b/crates/pilotty-cli/src/main.rs @@ -5,6 +5,7 @@ mod daemon; use clap::Parser; use pilotty_core::protocol::{Command, Request, ResponseData, ScrollDirection, SnapshotFormat}; +use std::io::Write; use tracing::{error, info}; use uuid::Uuid; @@ -50,6 +51,7 @@ fn cli_to_command(cli: &Cli) -> Option { .ok() .map(|p| p.to_string_lossy().into_owned()) }), + retain_bytes: args.retain_bytes, }), Commands::Kill(args) => Some(Command::Kill { session: args.session.clone(), @@ -88,6 +90,9 @@ fn cli_to_command(cli: &Cli) -> Option { session: args.session.clone(), }), Commands::ListSessions => Some(Command::ListSessions), + Commands::Logs(args) => Some(Command::Logs { + session: args.session.clone(), + }), Commands::Resize(args) => Some(Command::Resize { cols: args.cols, rows: args.rows, @@ -138,6 +143,20 @@ fn run_client_command(cli: Cli) -> anyhow::Result<()> { } => { println!("{}", content); } + ResponseData::Logs { + bytes, + total_bytes, + retained_bytes, + dropped_bytes, + truncated, + } => { + std::io::stdout().write_all(&bytes)?; + std::io::stdout().flush()?; + eprintln!( + "retention: total_bytes={total_bytes} retained_bytes={retained_bytes} \ + dropped_bytes={dropped_bytes} truncated={truncated}" + ); + } _ => println!("{}", serde_json::to_string_pretty(&data)?), } } diff --git a/crates/pilotty-core/src/protocol.rs b/crates/pilotty-core/src/protocol.rs index 3fb2af4..c4a87ad 100644 --- a/crates/pilotty-core/src/protocol.rs +++ b/crates/pilotty-core/src/protocol.rs @@ -77,6 +77,10 @@ pub enum Command { /// must be an existing directory. If not provided by the client, the /// process inherits the daemon's working directory. cwd: Option, + /// Maximum raw output bytes retained for this session. + /// Uses the daemon default when omitted. + #[serde(default)] + retain_bytes: Option, }, /// Kill a session. Kill { session: Option }, @@ -127,6 +131,8 @@ pub enum Command { }, /// List all active sessions. ListSessions, + /// Get the retained raw output for a session. + Logs { session: Option }, /// Resize the terminal. Resize { cols: u16, @@ -151,7 +157,14 @@ impl Command { /// compatibility before it can compile. pub fn minimum_protocol(&self) -> u32 { match self { - Self::Spawn { .. } + Self::Spawn { + retain_bytes: Some(_), + .. + } + | Self::Logs { .. } => PROTOCOL_VERSION, + Self::Spawn { + retain_bytes: None, .. + } | Self::Kill { .. } | Self::Snapshot { .. } | Self::Type { .. } @@ -261,6 +274,14 @@ pub enum ResponseData { }, /// Generic success message. Ok { message: String }, + /// Bounded raw output retained for a session. + Logs { + bytes: Vec, + total_bytes: u64, + retained_bytes: u64, + dropped_bytes: u64, + truncated: bool, + }, } impl ResponseData { @@ -270,6 +291,7 @@ impl ResponseData { /// older client. pub fn minimum_protocol(&self) -> u32 { match self { + Self::Logs { .. } => PROTOCOL_VERSION, Self::ScreenState(_) | Self::Snapshot { .. } | Self::SessionCreated { .. } @@ -384,4 +406,59 @@ mod tests { assert_eq!(command.minimum_protocol(), 0); assert_eq!(response.minimum_protocol(), 0); } + + #[test] + fn retention_commands_require_current_protocol() { + let plain_spawn = Command::Spawn { + command: vec!["sh".to_string()], + session_name: None, + cwd: None, + retain_bytes: None, + }; + let configured_spawn = Command::Spawn { + command: vec!["sh".to_string()], + session_name: None, + cwd: None, + retain_bytes: Some(1024), + }; + + assert_eq!(plain_spawn.minimum_protocol(), LEGACY_PROTOCOL_VERSION); + assert_eq!(configured_spawn.minimum_protocol(), PROTOCOL_VERSION); + assert_eq!( + Command::Logs { session: None }.minimum_protocol(), + PROTOCOL_VERSION + ); + } + + #[test] + fn logs_response_requires_current_protocol_and_preserves_raw_bytes() { + let response = ResponseData::Logs { + bytes: vec![0, 27, 255], + total_bytes: 9, + retained_bytes: 3, + dropped_bytes: 6, + truncated: true, + }; + + assert_eq!(response.minimum_protocol(), PROTOCOL_VERSION); + let json = serde_json::to_string(&response).expect("serialize logs response"); + let decoded: ResponseData = serde_json::from_str(&json).expect("deserialize logs response"); + assert_eq!(decoded, response); + } + + #[test] + fn legacy_spawn_without_retention_field_uses_daemon_default() { + let json = r#"{"action":"spawn","command":["sh"],"session_name":null,"cwd":null}"#; + let command: Command = serde_json::from_str(json).expect("deserialize legacy spawn"); + + assert_eq!( + command, + Command::Spawn { + command: vec!["sh".to_string()], + session_name: None, + cwd: None, + retain_bytes: None, + } + ); + } }