| 1 | //! Pseudo-terminal session wrapping `portable-pty`. |
| 2 | //! |
| 3 | //! Spawns a binary in a real PTY, pumps the child's stdout into an in-memory |
| 4 | //! buffer on a background thread, and exposes write/wait/kill primitives |
| 5 | //! the test harness composes. |
| 6 | //! |
| 7 | //! The reader thread is necessary because `portable-pty`'s reader is blocking |
| 8 | //! and the test thread must remain free to send input + poll for screen |
| 9 | //! changes. |
| 10 | |
| 11 | use anyhow::{Context, Result}; |
| 12 | use portable_pty::{Child, CommandBuilder, MasterPty, PtySize, native_pty_system}; |
| 13 | use std::io::{Read, Write}; |
| 14 | use std::path::Path; |
| 15 | use std::sync::{Arc, Mutex}; |
| 16 | use std::thread::{self, JoinHandle}; |
| 17 | use std::time::{Duration, Instant}; |
| 18 | |
| 19 | pub struct PtySession { |
| 20 | /// Held (not read) so the PTY master stays open for the child's lifetime. |
| 21 | master: Box<dyn MasterPty + Send>, |
| 22 | child: Box<dyn Child + Send + Sync>, |
| 23 | writer: Box<dyn Write + Send>, |
| 24 | buffer: Arc<Mutex<Vec<u8>>>, |
| 25 | /// Every byte the child ever wrote, never drained. `buffer` is consumed by |
| 26 | /// the frame parser, which is the wrong shape for assertions about the |
| 27 | /// control stream itself — terminal-mode setup/teardown is only visible as |
| 28 | /// escape sequences, and a mode that was enabled and then disabled leaves |
| 29 | /// no trace on the rendered screen at all. |
| 30 | transcript: Arc<Mutex<Vec<u8>>>, |
| 31 | reader_handle: Option<JoinHandle<()>>, |
| 32 | } |
| 33 | |
| 34 | pub struct PtySessionBuilder<'a> { |
| 35 | program: &'a Path, |
| 36 | args: Vec<String>, |
| 37 | cwd: Option<&'a Path>, |
| 38 | env: Vec<(String, String)>, |
| 39 | rows: u16, |
| 40 | cols: u16, |
| 41 | clear_env: bool, |
| 42 | } |
| 43 | |
| 44 | impl<'a> PtySessionBuilder<'a> { |
| 45 | pub fn new(program: &'a Path) -> Self { |
| 46 | Self { |
| 47 | program, |
| 48 | args: Vec::new(), |
| 49 | cwd: None, |
| 50 | env: Vec::new(), |
| 51 | rows: 40, |
| 52 | cols: 120, |
| 53 | clear_env: false, |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | pub fn args<I, S>(mut self, args: I) -> Self |
| 58 | where |
| 59 | I: IntoIterator<Item = S>, |
| 60 | S: Into<String>, |
| 61 | { |
| 62 | self.args.extend(args.into_iter().map(Into::into)); |
| 63 | self |
| 64 | } |
| 65 | |
| 66 | pub fn cwd(mut self, p: &'a Path) -> Self { |
| 67 | self.cwd = Some(p); |
| 68 | self |
| 69 | } |
| 70 | |
| 71 | pub fn env(mut self, k: impl Into<String>, v: impl Into<String>) -> Self { |
| 72 | self.env.push((k.into(), v.into())); |
| 73 | self |
| 74 | } |
| 75 | |
| 76 | /// Wipe the inherited environment before applying explicit `env(..)` |
| 77 | /// overrides. Use for sealed scenarios that must not see the developer's |
| 78 | /// real `~/.deepseek/`, `$HOME`, or API keys. |
| 79 | pub fn clear_env(mut self, yes: bool) -> Self { |
| 80 | self.clear_env = yes; |
| 81 | self |
| 82 | } |
| 83 | |
| 84 | pub fn size(mut self, rows: u16, cols: u16) -> Self { |
| 85 | self.rows = rows; |
| 86 | self.cols = cols; |
| 87 | self |
| 88 | } |
| 89 | |
| 90 | pub fn spawn(self) -> Result<PtySession> { |
| 91 | let pty_system = native_pty_system(); |
| 92 | let pair = pty_system |
| 93 | .openpty(PtySize { |
| 94 | rows: self.rows, |
| 95 | cols: self.cols, |
| 96 | pixel_width: 0, |
| 97 | pixel_height: 0, |
| 98 | }) |
| 99 | .context("openpty")?; |
| 100 | |
| 101 | let mut cmd = CommandBuilder::new(self.program); |
| 102 | for a in &self.args { |
| 103 | cmd.arg(a); |
| 104 | } |
| 105 | if let Some(cwd) = self.cwd { |
| 106 | cmd.cwd(cwd); |
| 107 | } |
| 108 | if self.clear_env { |
| 109 | cmd.env_clear(); |
| 110 | if let Some(path) = std::env::var_os("PATH") { |
| 111 | cmd.env("PATH", path); |
| 112 | } |
| 113 | } |
| 114 | // TERM must be set to something xterm-ish so crossterm enables the |
| 115 | // capabilities the TUI assumes (256 color, bracketed paste, …). |
| 116 | cmd.env("TERM", "xterm-256color"); |
| 117 | cmd.env("COLORTERM", "truecolor"); |
| 118 | for (k, v) in &self.env { |
| 119 | cmd.env(k, v); |
| 120 | } |
| 121 | |
| 122 | let child = pair.slave.spawn_command(cmd).context("spawn child")?; |
| 123 | // Drop the slave end so EOF propagates correctly when the child exits. |
| 124 | drop(pair.slave); |
| 125 | |
| 126 | let mut reader = pair.master.try_clone_reader().context("clone reader")?; |
| 127 | let writer = pair.master.take_writer().context("take writer")?; |
| 128 | |
| 129 | let buffer: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new())); |
| 130 | let transcript: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new())); |
| 131 | let buf_thread = Arc::clone(&buffer); |
| 132 | let transcript_thread = Arc::clone(&transcript); |
| 133 | let reader_handle = thread::Builder::new() |
| 134 | .name("qa-pty-reader".into()) |
| 135 | .spawn(move || { |
| 136 | let mut chunk = [0u8; 8192]; |
| 137 | loop { |
| 138 | match reader.read(&mut chunk) { |
| 139 | Ok(0) => break, |
| 140 | Ok(n) => { |
| 141 | if let Ok(mut b) = buf_thread.lock() { |
| 142 | b.extend_from_slice(&chunk[..n]); |
| 143 | } |
| 144 | if let Ok(mut t) = transcript_thread.lock() { |
| 145 | t.extend_from_slice(&chunk[..n]); |
| 146 | } |
| 147 | } |
| 148 | Err(_) => break, |
| 149 | } |
| 150 | } |
| 151 | }) |
| 152 | .context("reader thread")?; |
| 153 | |
| 154 | Ok(PtySession { |
| 155 | master: pair.master, |
| 156 | child, |
| 157 | writer, |
| 158 | buffer, |
| 159 | transcript, |
| 160 | reader_handle: Some(reader_handle), |
| 161 | }) |
| 162 | } |
| 163 | } |
| 164 | |
| 165 | impl PtySession { |
| 166 | pub fn builder(program: &Path) -> PtySessionBuilder<'_> { |
| 167 | PtySessionBuilder::new(program) |
| 168 | } |
| 169 | |
| 170 | pub fn pid(&self) -> Option<u32> { |
| 171 | self.child.process_id() |
| 172 | } |
| 173 | |
| 174 | pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> { |
| 175 | self.writer.write_all(bytes).context("pty write")?; |
| 176 | self.writer.flush().context("pty flush")?; |
| 177 | Ok(()) |
| 178 | } |
| 179 | |
| 180 | pub fn resize(&self, rows: u16, cols: u16) -> Result<()> { |
| 181 | self.master |
| 182 | .resize(PtySize { |
| 183 | rows, |
| 184 | cols, |
| 185 | pixel_width: 0, |
| 186 | pixel_height: 0, |
| 187 | }) |
| 188 | .context("pty resize") |
| 189 | } |
| 190 | |
| 191 | /// Drain any bytes the reader thread has pushed into the buffer. Returns |
| 192 | /// the bytes read this call. Non-blocking — returns immediately even if |
| 193 | /// the buffer is empty. |
| 194 | /// Every byte the child has written so far, including bytes already fed |
| 195 | /// to the frame parser. Non-destructive, so it can be sampled repeatedly. |
| 196 | pub fn transcript(&self) -> Vec<u8> { |
| 197 | self.transcript |
| 198 | .lock() |
| 199 | .unwrap_or_else(|e| e.into_inner()) |
| 200 | .clone() |
| 201 | } |
| 202 | |
| 203 | pub fn drain(&mut self) -> Vec<u8> { |
| 204 | let mut b = self.buffer.lock().unwrap_or_else(|e| e.into_inner()); |
| 205 | std::mem::take(&mut *b) |
| 206 | } |
| 207 | |
| 208 | /// Block until the child exits or the deadline passes. Returns the exit |
| 209 | /// status if reaped, or `None` on timeout. |
| 210 | pub fn wait_until(&mut self, deadline: Instant) -> Option<i32> { |
| 211 | loop { |
| 212 | match self.child.try_wait() { |
| 213 | Ok(Some(status)) => return Some(status.exit_code() as i32), |
| 214 | Ok(None) => {} |
| 215 | Err(_) => return None, |
| 216 | } |
| 217 | if Instant::now() >= deadline { |
| 218 | return None; |
| 219 | } |
| 220 | thread::sleep(Duration::from_millis(20)); |
| 221 | } |
| 222 | } |
| 223 | |
| 224 | /// Send SIGTERM-equivalent and wait briefly. Returns the exit status if |
| 225 | /// the child reaped within `grace`, or `None` otherwise. |
| 226 | pub fn shutdown(mut self, grace: Duration) -> Option<i32> { |
| 227 | self.kill_and_join_reader(grace) |
| 228 | } |
| 229 | |
| 230 | fn kill_and_join_reader(&mut self, grace: Duration) -> Option<i32> { |
| 231 | let _ = self.child.kill(); |
| 232 | let exit = self.wait_until(Instant::now() + grace); |
| 233 | if exit.is_some() |
| 234 | && let Some(handle) = self.reader_handle.take() |
| 235 | { |
| 236 | // Don't block on the reader thread forever — it exits on EOF. |
| 237 | let _ = handle.join(); |
| 238 | } |
| 239 | exit |
| 240 | } |
| 241 | } |
| 242 | |
| 243 | impl Drop for PtySession { |
| 244 | fn drop(&mut self) { |
| 245 | let _ = self.kill_and_join_reader(Duration::from_secs(2)); |
| 246 | } |
| 247 | } |
| 248 |