winx_code_agent/state/pty.rs
1//! Real PTY implementation using portable-pty
2//!
3//! This module provides a true pseudo-terminal interface for interactive
4//! shell sessions, enabling proper handling of:
5//! - ANSI escape sequences and colors
6//! - Interactive programs (sudo, vim, less, etc.)
7//! - Terminal resize events
8//! - Job control signals (Ctrl+C, Ctrl+Z, etc.)
9
10use anyhow::{anyhow, Context, Result};
11use portable_pty::{native_pty_system, Child, CommandBuilder, MasterPty, PtySize};
12use std::collections::hash_map::DefaultHasher;
13use std::hash::{Hash, Hasher};
14use std::io::{Read, Write};
15use std::path::{Path, PathBuf};
16use std::process::Command;
17use std::sync::mpsc::{self, TryRecvError};
18use std::sync::{Arc, Mutex as StdMutex};
19use std::thread;
20use std::time::{Duration, Instant};
21use tokio::sync::Mutex;
22use tracing::{debug, info, warn};
23
24use crate::state::line_ring::LineRing;
25use crate::state::live_terminal::{LiveTerminal, ScreenUpdate};
26
27/// Default terminal dimensions (columns x rows)
28pub const DEFAULT_COLS: u16 = 200;
29pub const DEFAULT_ROWS: u16 = 50;
30
31/// Maximum output buffer size to prevent memory issues
32const MAX_OUTPUT_SIZE: usize = 1_000_000;
33
34/// Cap on bytes streamed to one command's output-offload scratch file, so a
35/// command emitting gigabytes cannot fill the disk. The head beyond this is
36/// dropped (the agent still gets the first 50 MB plus the live tail).
37const SCRATCH_MAX_BYTES: u64 = 50 * 1024 * 1024;
38
39/// How many fully-formed lines to keep in the per-shell ringbuffer. Callers can
40/// ask for at most this many lines of historical context via
41/// `StatusCheck.scrollback_lines`.
42pub const RING_BUFFER_LINES: usize = 2_000;
43
44/// Cap on the unterminated tail held in `line_ring_partial`. Bounds memory when
45/// a program emits megabytes without ever printing a newline (a `\r`-redrawn
46/// progress bar, a binary blob). Generous enough to hold any real terminal line.
47const MAX_PARTIAL_LINE_BYTES: usize = 64 * 1024;
48
49/// WCGW-style prompt pattern for command completion detection
50const WCGW_PROMPT_PATTERN: &str = "◉";
51const WCGW_PROMPT_END: &str = "──➤";
52
53fn attachable_command(restricted_mode: bool) -> (CommandBuilder, Option<String>, bool) {
54 let requested = std::env::var("WINX_ATTACH_TERMINAL")
55 .or_else(|_| std::env::var("WINX_USE_SCREEN"))
56 .unwrap_or_default();
57 if !requested.is_empty() && requested != "0" && requested != "false" {
58 let session = format!("winx-{}-{}", std::process::id(), timestamp_millis());
59 if requested == "tmux" && command_available("tmux") {
60 let mut cmd = CommandBuilder::new("tmux");
61 cmd.args(["new-session", "-A", "-s", &session, "bash"]);
62 if restricted_mode {
63 cmd.arg("-r");
64 }
65 return (cmd, Some(format!("tmux attach -t {session}")), false);
66 }
67 if command_available("screen") {
68 // Parity with wcgw: ensure a sane ~/.screenrc and reap sessions whose
69 // creating winx process has died before spawning a fresh one.
70 ensure_screenrc();
71 cleanup_orphaned_screens();
72 let mut cmd = CommandBuilder::new("screen");
73 cmd.args(["-q", "-S", &session, "bash"]);
74 if restricted_mode {
75 cmd.arg("-r");
76 }
77 return (cmd, Some(format!("screen -x {session}")), false);
78 }
79 }
80
81 let shell = preferred_shell(restricted_mode);
82 let is_zsh = shell == "zsh";
83 let mut cmd = CommandBuilder::new(&shell);
84 // zsh's restricted mode isn't the `-r` flag, so restricted always uses bash -r.
85 if restricted_mode && !is_zsh {
86 cmd.arg("-r");
87 }
88 (cmd, None, is_zsh)
89}
90
91/// Shell to spawn directly. Defaults to bash; honors `WINX_SHELL=zsh` when zsh is
92/// on PATH and we're not in restricted mode (zsh's restricted mode differs from
93/// `bash -r`, so restricted falls back to bash).
94fn preferred_shell(restricted_mode: bool) -> String {
95 if !restricted_mode {
96 if let Ok(requested) = std::env::var("WINX_SHELL") {
97 if requested == "zsh" && command_available("zsh") {
98 return "zsh".to_string();
99 }
100 }
101 }
102 "bash".to_string()
103}
104
105fn command_available(command: &str) -> bool {
106 Command::new("sh")
107 .args(["-c", &format!("command -v {command}")])
108 .output()
109 .is_ok_and(|output| output.status.success())
110}
111
112/// Create `~/.screenrc` with a large scrollback if the user has none, matching
113/// wcgw's `check_if_screen_command_available`. Never overwrites an existing file.
114fn ensure_screenrc() {
115 let Some(home) = home::home_dir() else {
116 return;
117 };
118 let screenrc = home.join(".screenrc");
119 if screenrc.exists() {
120 return;
121 }
122 let _ = std::fs::write(
123 &screenrc,
124 "defscrollback 10000\ntermcapinfo xterm* ti@:te@\nstartup_message off\n",
125 );
126}
127
128/// Reap detached `winx-*` screen sessions whose creating process is gone.
129///
130/// The session name embeds the creator PID (`winx-<pid>-<ts>`), so an orphan is
131/// simply a session whose `<pid>` no longer exists — the wcgw equivalent of
132/// detecting `parent_pid == 1`. Best-effort: any failure is silently ignored.
133fn cleanup_orphaned_screens() {
134 let Ok(output) = Command::new("screen").arg("-ls").output() else {
135 return;
136 };
137 // `screen -ls` exits non-zero when sessions exist, so we parse stdout regardless.
138 let listing = String::from_utf8_lossy(&output.stdout);
139 for line in listing.lines() {
140 let Some(session) = line.split_whitespace().next() else {
141 continue;
142 };
143 // session token looks like "<screen_pid>.winx-<creator_pid>-<ts>"
144 let Some((_, name)) = session.split_once('.') else {
145 continue;
146 };
147 if let Some(creator_pid) = winx_creator_pid(name) {
148 if !process_exists(creator_pid) {
149 let _ = Command::new("screen").args(["-S", session, "-X", "quit"]).output();
150 }
151 }
152 }
153}
154
155/// Extract the creator PID from a `winx-<pid>-<ts>` screen session name.
156fn winx_creator_pid(name: &str) -> Option<u32> {
157 name.strip_prefix("winx-")?.split('-').next()?.parse::<u32>().ok()
158}
159
160/// Whether a process with `pid` is currently alive (Linux `/proc` check).
161fn process_exists(pid: u32) -> bool {
162 std::path::Path::new("/proc").join(pid.to_string()).exists()
163}
164
165fn timestamp_millis() -> u128 {
166 std::time::SystemTime::now()
167 .duration_since(std::time::UNIX_EPOCH)
168 .map_or(0, |duration| duration.as_millis())
169}
170
171/// Real PTY-based interactive shell
172///
173/// Uses portable-pty for true pseudo-terminal functionality,
174/// enabling proper handling of interactive programs like sudo, vim, etc.
175pub struct PtyShell {
176 /// The PTY master handle for resize operations
177 master: Box<dyn MasterPty + Send>,
178 /// Child process running the shell
179 child: Box<dyn Child + Send + Sync>,
180 /// Writer for PTY input (taken from master)
181 writer: Box<dyn Write + Send>,
182 /// Channel receiver for output from reader thread
183 output_rx: mpsc::Receiver<String>,
184 /// Current terminal size
185 size: PtySize,
186 /// Last command executed
187 pub last_command: String,
188 /// Accumulated output buffer
189 pub output_buffer: String,
190 /// Whether a command is currently running
191 pub command_running: bool,
192 /// Last cwd observed in this shell's nonce-bearing prompt.
193 current_cwd: PathBuf,
194 /// Monotonic start time of the current or most recently completed command.
195 /// Status polling uses this instead of restarting the clock per tool call.
196 command_started_at: Option<Instant>,
197 /// Exit code of the last completed foreground command, parsed from the prompt
198 /// marker (`──➤<nonce>:<code>`). `None` until a command finishes, and reset
199 /// to `None` while one is running.
200 pub last_exit_code: Option<i32>,
201 /// Maximum output size before truncation
202 max_output_size: usize,
203 /// Flag for output truncation
204 pub output_truncated: bool,
205 /// Bounded scrollback ring of raw PTY lines (newest at the back) plus the
206 /// in-flight partial line, with drop accounting so a truncated scrollback is
207 /// reported to the caller instead of silently shrinking. See `LineRing`.
208 ring: LineRing,
209 /// Hash of the last rendered output we shipped to the caller. Used by the
210 /// delta path in `status_check` to elide repeats when the screen is idle.
211 pub last_returned_hash: Option<u64>,
212 /// Cumulative PTY snapshot already delivered to the caller. Incremental
213 /// status polling diffs against this value, then advances it after rendering.
214 delivered_output: String,
215 /// Optional command a human can run to attach to the same terminal session.
216 pub attach_hint: Option<String>,
217 /// The exact suffix that ends this shell's prompt: `──➤<nonce>`. The nonce is
218 /// a per-shell random value embedded in `PROMPT_COMMAND`, so command output
219 /// can't impersonate the prompt (printing `◉ x──➤` no longer ends a command
220 /// early). Anchoring completion on this instead of the bare glyphs removes
221 /// the false-positive that truncated output when a program echoed them.
222 prompt_end_marker: String,
223 /// Live terminal emulator (fixed-viewport, `vt100`-backed) fed continuously
224 /// by the reader thread. Unlike the scrollback ring (raw lines), it keeps a
225 /// consolidated screen grid — cursor moves, redraws, alternate-screen
226 /// applied — so a TUI yields a stable, non-stacked snapshot. The fixed
227 /// viewport is what lets inline-rendering TUIs (the `agy`/Antigravity CLI,
228 /// which redraws with cursor-up + erase-line in the main screen) consolidate
229 /// instead of ghosting. See `live_snapshot`.
230 live: Arc<StdMutex<LiveTerminal>>,
231 /// Workspace root under which output-offload scratch files are written.
232 /// `None` until a caller sets it (the workspace is only known after
233 /// Initialize, and `execute_command` sets it per command). See
234 /// `crate::utils::scratch_file`.
235 scratch_workspace_root: Option<PathBuf>,
236 /// Scratch file holding the CURRENT command's dropped output head, created
237 /// lazily on the first truncation and reset per command.
238 scratch_path: Option<PathBuf>,
239 /// Bytes already streamed to `scratch_path`, used to enforce `SCRATCH_MAX_BYTES`.
240 scratch_bytes: u64,
241}
242
243impl std::fmt::Debug for PtyShell {
244 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
245 f.debug_struct("PtyShell")
246 .field("size", &format!("{}x{}", self.size.cols, self.size.rows))
247 .field("last_command", &self.last_command)
248 .field("command_running", &self.command_running)
249 .field("current_cwd", &self.current_cwd)
250 .field("output_truncated", &self.output_truncated)
251 .field("output_buffer_len", &self.output_buffer.len())
252 .field("attach_hint", &self.attach_hint)
253 .finish_non_exhaustive()
254 }
255}
256
257/// Resolve the process group to signal for a PTY child, but ONLY when it is
258/// safe: the pid must be a real child (`> 1`) that leads its own group
259/// (`getpgid(pid) == pid`). portable-pty `setsid`s the shell, so a live PTY
260/// child is always its own group leader; if that ever stopped holding,
261/// `kill(-pgid, ...)` would hit some unrelated group, so we refuse instead of
262/// guessing.
263///
264/// The `pid > 1` check is load-bearing: `kill(-1, ...)` means "every process I
265/// may signal" and `kill(0, ...)` means "my own group" - either would be
266/// catastrophic from a teardown path.
267#[cfg(unix)]
268fn killable_group(pid: u32) -> Option<i32> {
269 let pid = i32::try_from(pid).ok()?;
270 if pid <= 1 {
271 return None;
272 }
273 // SAFETY: getpgid(2) reads the group of an existing pid; integer-only, no
274 // memory touched. Returns -1 for a dead/invalid pid, which fails the check.
275 let group = unsafe { libc::getpgid(pid) };
276 (group == pid).then_some(group)
277}
278
279/// Send `signal` to the whole process group `pgid` (group-kill via negative
280/// pid). `pgid` must be one vetted by [`killable_group`].
281#[cfg(unix)]
282fn signal_group(pgid: i32, signal: i32) {
283 // SAFETY: kill(2) with a negative pid signals a process group; integer-only,
284 // no memory touched. `pgid` came from `killable_group`, so it is `> 1`.
285 unsafe {
286 libc::kill(-pgid, signal);
287 }
288}
289
290impl Drop for PtyShell {
291 /// Kill and reap the shell child **and its whole process group** so neither
292 /// the shell nor anything it spawned leaks.
293 ///
294 /// `std::process::Child::drop` neither kills nor waits, and even
295 /// `Child::kill` signals only the shell itself: a background job it started
296 /// (`npm run dev &`, a `nohup`'d server) would be reparented to init and live
297 /// on. Because portable-pty runs the shell as a session/group leader, we can
298 /// signal the whole group with a negative pid - SIGTERM first so well-behaved
299 /// children can clean up, then `Child::kill` (SIGKILL on the leader, which
300 /// closes the PTY slave so the reader thread's `read()` returns EOF and the
301 /// thread exits), then SIGKILL the group to sweep anything still alive, then
302 /// reap. The pgid is captured while the leader is alive and reused: it stays
303 /// valid until we `wait()`, even after the leader dies. All best-effort.
304 fn drop(&mut self) {
305 // Only group-kill the direct bash/zsh path. Under screen/tmux (`attach_hint`
306 // is `Some`) the `Child` is a multiplexer *client*; the real shell runs in
307 // a detached daemon in another process group, so a group-kill here would
308 // miss the shell entirely and could disturb unrelated sessions of the same
309 // multiplexer. Let the multiplexer own its lifecycle in that case.
310 #[cfg(unix)]
311 let pgid = if self.attach_hint.is_some() {
312 None
313 } else {
314 self.child.process_id().and_then(killable_group)
315 };
316
317 #[cfg(unix)]
318 if let Some(pgid) = pgid {
319 signal_group(pgid, libc::SIGTERM);
320 }
321
322 // SIGKILL the leader and close the PTY slave (unblocks the reader thread).
323 let _ = self.child.kill();
324
325 #[cfg(unix)]
326 if let Some(pgid) = pgid {
327 signal_group(pgid, libc::SIGKILL);
328 }
329
330 let _ = self.child.wait();
331 }
332}
333
334impl PtyShell {
335 /// Create a new PTY shell session
336 ///
337 /// # Arguments
338 /// * `initial_dir` - Starting directory for the shell
339 /// * `restricted_mode` - Whether to use bash restricted mode (-r)
340 ///
341 /// # Returns
342 /// A new `PtyShell` instance with an active bash session
343 pub fn new(initial_dir: &Path, restricted_mode: bool) -> Result<Self> {
344 info!(
345 "Creating new PTY shell (restricted: {}) in {}",
346 restricted_mode,
347 initial_dir.display()
348 );
349
350 // Initialize the native PTY system
351 let pty_system = native_pty_system();
352
353 // Configure terminal size
354 let size =
355 PtySize { rows: DEFAULT_ROWS, cols: DEFAULT_COLS, pixel_width: 0, pixel_height: 0 };
356
357 // Open the PTY pair (master + slave)
358 let pair = pty_system.openpty(size).context("Failed to open PTY pair")?;
359
360 // Build the command
361 let (mut cmd, attach_hint, is_zsh) = attachable_command(restricted_mode);
362
363 // Set up environment for proper terminal behavior
364 cmd.env("TERM", "xterm-256color");
365 cmd.env("COLORTERM", "truecolor");
366 cmd.env("PAGER", "cat");
367 cmd.env("GIT_PAGER", "cat");
368 cmd.env("COLUMNS", DEFAULT_COLS.to_string());
369 cmd.env("ROWS", DEFAULT_ROWS.to_string());
370 // Per-shell random nonce embedded in the prompt so command output can't
371 // forge the completion marker. 64 bits of entropy in hex.
372 let nonce = format!("{:016x}", rand::random::<u64>());
373 // WCGW-style prompt for command completion detection
374 // Note: removed \r\e[2K which was erasing the prompt before it could be detected
375 // `__winx_ec=$?` captures the user command's exit status as the very
376 // first thing PROMPT_COMMAND does, then we print it after the marker.
377 cmd.env(
378 "PROMPT_COMMAND",
379 format!(r#"__winx_ec=$?; printf "◉ %s──➤{nonce}:%s " "$PWD" "$__winx_ec""#),
380 );
381 cmd.cwd(initial_dir);
382
383 // Spawn bash in the PTY slave
384 let child = pair.slave.spawn_command(cmd).context("Failed to spawn bash in PTY")?;
385
386 // Get reader and writer from master
387 let mut reader = pair.master.try_clone_reader().context("Failed to clone PTY reader")?;
388 let writer = pair.master.take_writer().context("Failed to take PTY writer")?;
389
390 // Bounded channel: if the consumer stalls, the reader thread blocks on
391 // send (natural backpressure via the PTY OS buffer) instead of growing
392 // memory without bound on a runaway producer like `yes`.
393 let (output_tx, output_rx) = mpsc::sync_channel::<String>(1024);
394
395 // Live terminal emulator, shared with the reader thread so the screen
396 // grid stays current without any consumer needing to poll.
397 let live = Arc::new(StdMutex::new(LiveTerminal::new(DEFAULT_ROWS, DEFAULT_COLS)));
398 let live_reader = Arc::clone(&live);
399
400 // Spawn a background thread to read from the PTY
401 // This prevents blocking the main thread
402 thread::spawn(move || {
403 let mut buf = [0u8; 4096];
404 // Bytes of an incomplete trailing UTF-8 char held back for the next
405 // read. Without this, a multibyte glyph split across a 4096-byte read
406 // boundary would decode to two U+FFFDs — corrupting the output buffer
407 // and the prompt/CWD detection that runs over it.
408 let mut carry: Vec<u8> = Vec::new();
409 loop {
410 match reader.read(&mut buf) {
411 Ok(0) => {
412 // EOF - PTY closed. Flush any held bytes lossily.
413 if !carry.is_empty() {
414 let _ = output_tx.send(String::from_utf8_lossy(&carry).into_owned());
415 }
416 break;
417 }
418 Ok(n) => {
419 // Tap the raw bytes into the live emulator first (brief
420 // lock; feed is O(chunk len)). Feeding bytes — not the
421 // lossy String — keeps the persistent VTE parser exact
422 // across chunk boundaries. Recover from a poisoned lock
423 // instead of dropping the feed silently.
424 {
425 let mut emu = live_reader
426 .lock()
427 .unwrap_or_else(std::sync::PoisonError::into_inner);
428 emu.feed(&buf[..n]);
429 }
430 // Decode the longest valid prefix; carry an incomplete
431 // trailing char into the next read.
432 carry.extend_from_slice(&buf[..n]);
433 let (chunk, rest) = decode_keep_incomplete(&carry);
434 carry = rest;
435 if !chunk.is_empty() && output_tx.send(chunk).is_err() {
436 // Receiver dropped, exit thread
437 break;
438 }
439 }
440 Err(e) => {
441 debug!("PTY reader thread error: {}", e);
442 break;
443 }
444 }
445 }
446 debug!("PTY reader thread exiting");
447 });
448
449 // Create the shell instance
450 let mut shell = Self {
451 master: pair.master,
452 child,
453 writer,
454 output_rx,
455 size,
456 last_command: String::new(),
457 output_buffer: String::new(),
458 command_running: false,
459 current_cwd: initial_dir.to_path_buf(),
460 command_started_at: None,
461 last_exit_code: None,
462 max_output_size: MAX_OUTPUT_SIZE,
463 output_truncated: false,
464 ring: LineRing::new(RING_BUFFER_LINES, MAX_PARTIAL_LINE_BYTES),
465 last_returned_hash: None,
466 delivered_output: String::new(),
467 attach_hint,
468 prompt_end_marker: format!("{WCGW_PROMPT_END}{nonce}"),
469 live,
470 scratch_workspace_root: None,
471 scratch_path: None,
472 scratch_bytes: 0,
473 };
474
475 // Initialize the shell with WCGW-style prompt
476 shell.initialize_prompt(is_zsh, &nonce)?;
477
478 debug!("PTY shell created successfully");
479 Ok(shell)
480 }
481
482 /// Initialize the shell prompt for WCGW compatibility
483 fn initialize_prompt(&mut self, is_zsh: bool, nonce: &str) -> Result<()> {
484 // Set up the dynamic prompt - matches WCGW Python PROMPT_STATEMENT.
485 // zsh ignores PROMPT_COMMAND, so it gets a precmd hook with a blanked
486 // default prompt instead.
487 // Blank the shell's own PS1/PROMPT so the line ends exactly at our `──➤`
488 // marker. Otherwise a user's ~/.bashrc/~/.zshrc PS1 (e.g. `[user@host]$`)
489 // is appended after the marker, and prompt detection — which anchors on a
490 // trailing `──➤` — only fires when the chunk happens to fragment right
491 // before the PS1, making command-completion detection flaky.
492 let prompt_statement = if is_zsh {
493 // Clear precmd_functions too: frameworks (oh-my-zsh/p10k) register the
494 // prompt via that array, so redefining `precmd` alone leaves their
495 // prompt in place and our `──➤` marker never ends the line.
496 format!(
497 r#"export GIT_PAGER=cat PAGER=cat; precmd_functions=(); preexec_functions=(); PROMPT=''; RPROMPT=''; precmd() {{ local __winx_ec=$?; printf "◉ %s──➤{nonce}:%s " "$PWD" "$__winx_ec" }}"#
498 )
499 } else {
500 format!(
501 r#"export GIT_PAGER=cat PAGER=cat PROMPT_COMMAND='__winx_ec=$?; printf "◉ %s──➤{nonce}:%s " "$PWD" "$__winx_ec"'; PS1=''"#
502 )
503 };
504
505 self.write_command(&prompt_statement)?;
506
507 // Wait for prompt to be ready
508 std::thread::sleep(Duration::from_millis(100));
509 let _ = self.drain_output();
510
511 Ok(())
512 }
513
514 /// Write a command to the PTY, submitting it with a carriage return.
515 fn write_command(&mut self, command: &str) -> Result<()> {
516 // Submit with `\r`, matching the foreground path (`send_special_key("Enter")`
517 // also sends `\r`) and a real terminal's Enter key. The two paths used to
518 // disagree — background sent `\n` here, foreground sent `\r` — which is a
519 // real difference for readline-based TUIs. bash's canonical line discipline
520 // maps CR to NL (ICRNL), so plain commands behave identically.
521 let cmd_with_newline = format!("{command}\r");
522 self.writer.write_all(cmd_with_newline.as_bytes()).context("Failed to write to PTY")?;
523 self.writer.flush().context("Failed to flush PTY")?;
524 Ok(())
525 }
526
527 /// Drain any pending output from the PTY channel
528 fn drain_output(&mut self) -> String {
529 let mut output = String::new();
530 let deadline = Instant::now() + Duration::from_millis(200);
531
532 // Drain all available output from the channel
533 while Instant::now() < deadline {
534 match self.output_rx.try_recv() {
535 Ok(chunk) => {
536 output.push_str(&chunk);
537
538 // Prevent runaway reads
539 if output.len() > self.max_output_size {
540 self.output_truncated = true;
541 break;
542 }
543 }
544 Err(TryRecvError::Empty) => {
545 // No more data, wait briefly for more
546 thread::sleep(Duration::from_millis(10));
547 }
548 Err(TryRecvError::Disconnected) => {
549 // Reader thread died
550 break;
551 }
552 }
553 }
554
555 output
556 }
557
558 /// Send a command to the shell and start reading output
559 pub fn send_command(&mut self, command: &str) -> Result<()> {
560 debug!(bytes = command.len(), "PTY sending command");
561
562 // Clear previous state
563 self.output_buffer.clear();
564 self.output_truncated = false;
565 self.reset_scratch();
566 self.last_command = command.to_string();
567 self.command_running = true;
568 self.mark_command_started();
569 self.last_exit_code = None;
570 // A new command means the next status_check should return whatever
571 // shows up — drop the dedup hash so we don't elide the first response.
572 self.last_returned_hash = None;
573 self.delivered_output.clear();
574
575 // Write the command
576 self.write_command(command)?;
577
578 Ok(())
579 }
580
581 /// Start the monotonic runtime clock for a newly submitted command.
582 pub fn mark_command_started(&mut self) {
583 self.command_started_at = Some(Instant::now());
584 }
585
586 /// Elapsed runtime since the current command was submitted.
587 pub fn command_elapsed(&self) -> Option<Duration> {
588 self.command_started_at.map(|started_at| started_at.elapsed())
589 }
590
591 /// Current working directory last reported by this shell's prompt.
592 pub fn current_cwd(&self) -> &Path {
593 &self.current_cwd
594 }
595
596 /// Snapshot used as the left-hand side of incremental output rendering.
597 pub fn delivered_output(&self) -> String {
598 self.delivered_output.clone()
599 }
600
601 /// Mark the cumulative output snapshot as delivered after rendering it.
602 pub fn mark_output_delivered(&mut self, output: &str) {
603 self.delivered_output.clear();
604 self.delivered_output.push_str(output);
605 }
606
607 /// Return up to `lines` recent lines from the scrollback ring, oldest
608 /// first. When older lines were evicted or a long unterminated line was
609 /// clipped, a `[winx: …]` notice is prepended so the caller knows the
610 /// scrollback is incomplete rather than silently receiving less than asked.
611 pub fn collect_scrollback(&self, lines: usize) -> String {
612 // The ring holds raw PTY lines. Replay them through the terminal
613 // emulator so cursor movements (readline echo, in-place redraws) are
614 // *applied* — not merely stripped — yielding what the screen actually
615 // showed. A plain strip can't undo a `\x1b[D`, so it would leave
616 // `>>> p>>> pr>>> pri...` echo noise behind; the emulator collapses it.
617 let rendered =
618 crate::state::terminal::render_terminal_output(&self.ring.raw(lines)).join("\n");
619 match self.ring.scrollback_notice(lines) {
620 Some(notice) => format!("{notice}\n{rendered}"),
621 None => rendered,
622 }
623 }
624
625 /// Hash arbitrary rendered output into a u64 dedup key.
626 pub fn fingerprint(text: &str) -> u64 {
627 let mut hasher = DefaultHasher::new();
628 text.hash(&mut hasher);
629 hasher.finish()
630 }
631
632 /// Read output from the PTY with timeout
633 ///
634 /// Returns (output, `is_complete`) tuple where `is_complete` indicates
635 /// whether the command has finished (prompt detected)
636 pub fn read_output(&mut self, timeout_secs: f32) -> Result<(String, bool)> {
637 let timeout = Duration::from_secs_f32(timeout_secs.clamp(0.1, 60.0));
638 let start = Instant::now();
639 let mut complete = false;
640 let mut no_data_count = 0;
641 let mut prompt_detected_at: Option<Instant> = None;
642
643 while start.elapsed() < timeout {
644 match self.output_rx.try_recv() {
645 Ok(chunk) => {
646 self.output_buffer.push_str(&chunk);
647 self.ring.push_chunk(&chunk);
648 no_data_count = 0;
649
650 // Check for WCGW prompt indicating command completion
651 if prompt_detected_at.is_none()
652 && (Self::check_prompt_complete(&chunk, &self.prompt_end_marker)
653 || Self::check_prompt_complete(
654 &self.output_buffer,
655 &self.prompt_end_marker,
656 ))
657 {
658 prompt_detected_at = Some(Instant::now());
659 debug!("Prompt detected, draining remaining output...");
660 }
661
662 // Truncate if too large (offloads the dropped head to scratch).
663 self.truncate_output_buffer_with_offload();
664 }
665 Err(TryRecvError::Empty) => {
666 // No data available, wait briefly
667 thread::sleep(Duration::from_millis(10));
668 no_data_count += 1;
669
670 // If prompt was detected, check if we've drained long enough
671 if let Some(detected_time) = prompt_detected_at {
672 // Wait 100ms after prompt detection to capture any trailing output
673 if detected_time.elapsed() > Duration::from_millis(100) {
674 complete = true;
675 debug!("Command completed - prompt detected and drained");
676 break;
677 }
678 } else if no_data_count > 10
679 && Self::check_prompt_complete(&self.output_buffer, &self.prompt_end_marker)
680 {
681 // Prompt detected during empty reads
682 prompt_detected_at = Some(Instant::now());
683 debug!("Prompt detected after wait, draining...");
684 }
685 }
686 Err(TryRecvError::Disconnected) => {
687 // Reader thread died - PTY closed
688 warn!("PTY reader disconnected");
689 complete = true;
690 break;
691 }
692 }
693 }
694
695 if complete || prompt_detected_at.is_some() {
696 self.command_running = false;
697 self.last_exit_code =
698 Self::parse_prompt_exit_code(&self.output_buffer, &self.prompt_end_marker);
699 self.update_current_cwd_from_prompt();
700 complete = true;
701 }
702
703 Ok((self.output_buffer.clone(), complete))
704 }
705
706 /// Non-blocking drain of whatever the reader thread has already queued.
707 ///
708 /// Unlike [`read_output`](Self::read_output), this never sleeps: it consumes
709 /// every chunk currently available, updates prompt/completion state, and
710 /// returns immediately. `prune_finished_shells` uses it so the global
711 /// background-shell lock is never held across a blocking 100ms `read_output`.
712 /// Returns `true` if the command has completed (prompt seen or PTY closed).
713 pub fn poll_output_nonblocking(&mut self) -> bool {
714 let mut prompt_seen = false;
715 loop {
716 match self.output_rx.try_recv() {
717 Ok(chunk) => {
718 let chunk_has_prompt =
719 Self::check_prompt_complete(&chunk, &self.prompt_end_marker);
720 self.output_buffer.push_str(&chunk);
721 self.ring.push_chunk(&chunk);
722 if chunk_has_prompt
723 || Self::check_prompt_complete(&self.output_buffer, &self.prompt_end_marker)
724 {
725 prompt_seen = true;
726 }
727 self.truncate_output_buffer_with_offload();
728 }
729 Err(TryRecvError::Empty) => break,
730 Err(TryRecvError::Disconnected) => {
731 self.command_running = false;
732 return true;
733 }
734 }
735 }
736 if prompt_seen {
737 self.command_running = false;
738 self.last_exit_code =
739 Self::parse_prompt_exit_code(&self.output_buffer, &self.prompt_end_marker);
740 self.update_current_cwd_from_prompt();
741 }
742 // Completion is state, not an edge. Another consumer (notably the
743 // background reaper) may have drained the prompt first; a later poll
744 // must still report the already-idle shell as complete even when this
745 // particular drain saw no new prompt bytes.
746 !self.command_running
747 }
748
749 fn update_current_cwd_from_prompt(&mut self) {
750 let Some(last) =
751 self.output_buffer.lines().rev().find(|line| line.contains(&self.prompt_end_marker))
752 else {
753 return;
754 };
755 let clean = crate::state::terminal::strip_ansi_codes(last);
756 let Some((before_marker, _)) = clean.rsplit_once(&self.prompt_end_marker) else {
757 return;
758 };
759 let Some((_, cwd)) = before_marker.rsplit_once("◉ ") else {
760 return;
761 };
762 let cwd = cwd.trim();
763 if !cwd.is_empty() {
764 self.current_cwd = PathBuf::from(cwd);
765 }
766 }
767
768 /// Snapshot of the accumulated output buffer — the same text `read_output`
769 /// returns. Lets callers poll non-blocking (`poll_output_nonblocking`) and
770 /// grab the buffer once at the end instead of cloning it on every slice.
771 pub fn output_snapshot(&self) -> String {
772 self.output_buffer.clone()
773 }
774
775 /// Check if the output ends with this shell's prompt.
776 ///
777 /// `prompt_end` is the per-shell `──➤<nonce>` suffix. Anchoring on the LAST
778 /// non-empty line (not a global `contains`) plus the random nonce means
779 /// command output that happens to print the prompt glyphs mid-stream — or
780 /// even a forged `◉ x──➤ ` at end of line — can't be mistaken for the prompt
781 /// and truncate output or end the command early.
782 fn check_prompt_complete(text: &str, prompt_end: &str) -> bool {
783 Self::parse_prompt_exit_code(text, prompt_end).is_some()
784 }
785
786 /// If the last non-empty line is this shell's completion prompt
787 /// (`◉ <pwd>──➤<nonce>:<code>`), return the exit code that follows the
788 /// marker. Returns `None` when the prompt isn't present — which is exactly
789 /// "the command is still running". The per-shell nonce plus the structured
790 /// `:<digits>` suffix keep command output from forging completion.
791 fn parse_prompt_exit_code(text: &str, prompt_end: &str) -> Option<i32> {
792 let last = text.lines().rev().find(|line| !line.trim().is_empty())?;
793 // Strip ANSI so a trailing erase/cursor sequence after the code
794 // (e.g. ":0 \x1b[K") doesn't defeat the suffix check.
795 let clean = crate::state::terminal::strip_ansi_codes(last);
796 let clean = clean.trim_end();
797 if !clean.contains(WCGW_PROMPT_PATTERN) {
798 return None;
799 }
800 // Everything after the last "──➤<nonce>" must be ":<digits>" and nothing
801 // else — this confirms the prompt ends the line AND yields the code.
802 let after = clean.rsplit_once(prompt_end)?.1;
803 after.strip_prefix(':')?.trim().parse::<i32>().ok()
804 }
805
806 /// Point output-offload scratch files at `root` (the workspace). Idempotent;
807 /// callers set it per command so it always tracks the current workspace.
808 pub fn set_scratch_root(&mut self, root: &Path) {
809 self.scratch_workspace_root = Some(root.to_path_buf());
810 }
811
812 /// The scratch file holding the current command's dropped output head, if any
813 /// was offloaded this command.
814 pub fn scratch_path(&self) -> Option<&Path> {
815 self.scratch_path.as_deref()
816 }
817
818 /// Forget the current command's scratch file so the next command starts a
819 /// fresh one. Does not delete the file (pruned later by age).
820 pub fn reset_scratch(&mut self) {
821 self.scratch_path = None;
822 self.scratch_bytes = 0;
823 }
824
825 /// Append a just-dropped output `head` to the per-command scratch file so the
826 /// agent can recover it via `ReadFiles`. Best-effort: needs a known workspace
827 /// root, stops at `SCRATCH_MAX_BYTES`, and swallows IO errors.
828 fn offload_dropped_head(&mut self, head: &str) {
829 let Some(root) = self.scratch_workspace_root.clone() else {
830 return;
831 };
832 // Clamp to the remaining byte budget so the file honors SCRATCH_MAX_BYTES
833 // exactly (a single head can be up to keep_size, ~500 KB). Snap to a char
834 // boundary so we never slice mid-UTF-8.
835 let remaining = SCRATCH_MAX_BYTES.saturating_sub(self.scratch_bytes);
836 if remaining == 0 {
837 return;
838 }
839 let head = match usize::try_from(remaining) {
840 Ok(rem) if rem < head.len() => &head[..crate::utils::floor_char_boundary(head, rem)],
841 _ => head,
842 };
843 if self.scratch_path.is_none() {
844 self.scratch_path = crate::utils::scratch_file::new_scratch_path(&root);
845 }
846 let Some(path) = self.scratch_path.clone() else {
847 return;
848 };
849 match crate::utils::scratch_file::append_scratch(&path, head.as_bytes()) {
850 Ok(0) if self.scratch_bytes == 0 => self.scratch_path = None,
851 Ok(written) => {
852 self.scratch_bytes = self.scratch_bytes.saturating_add(written as u64);
853 }
854 Err(e) => debug!("scratch: append to {} failed: {e}", path.display()),
855 }
856 }
857
858 /// Once `output_buffer` exceeds the cap, offload the head about to be dropped
859 /// to the scratch file, then keep only the tail with a truncation marker.
860 /// Shared by the blocking and non-blocking drain paths so the offload and the
861 /// char-boundary handling live in one place.
862 fn truncate_output_buffer_with_offload(&mut self) {
863 if self.output_buffer.len() <= self.max_output_size {
864 return;
865 }
866 self.output_truncated = true;
867 let keep_size = self.max_output_size / 2;
868 // Snap to a char boundary: a raw byte offset can land mid-UTF-8
869 // (CJK/emoji/box-drawing/the prompt glyphs) and slicing there would panic.
870 let cut = crate::utils::floor_char_boundary(
871 &self.output_buffer,
872 self.output_buffer.len() - keep_size,
873 );
874 let head = self.output_buffer[..cut].to_owned();
875 self.offload_dropped_head(&head);
876 // Replace the dropped head in place instead of `format!`-ing a brand-new
877 // String: `replace_range` shifts the kept tail down once and reuses the
878 // buffer's allocation. `cut` is already a char boundary (above), and 0 is
879 // always one, so the range is valid.
880 self.output_buffer.replace_range(..cut, "\n(...output truncated...)\n");
881 }
882
883 /// Send Ctrl+C (interrupt) to the PTY
884 pub fn send_interrupt(&mut self) -> Result<()> {
885 debug!("PTY sending Ctrl+C");
886 self.writer
887 .write_all(&[0x03]) // ASCII ETX (Ctrl+C)
888 .context("Failed to send Ctrl+C")?;
889 self.writer.flush()?;
890 Ok(())
891 }
892
893 /// Send Ctrl+D (EOF) to the PTY
894 pub fn send_eof(&mut self) -> Result<()> {
895 debug!("PTY sending Ctrl+D");
896 self.writer
897 .write_all(&[0x04]) // ASCII EOT (Ctrl+D)
898 .context("Failed to send Ctrl+D")?;
899 self.writer.flush()?;
900 Ok(())
901 }
902
903 /// Send Ctrl+Z (suspend) to the PTY
904 pub fn send_suspend(&mut self) -> Result<()> {
905 debug!("PTY sending Ctrl+Z");
906 self.writer
907 .write_all(&[0x1A]) // ASCII SUB (Ctrl+Z)
908 .context("Failed to send Ctrl+Z")?;
909 self.writer.flush()?;
910 Ok(())
911 }
912
913 /// Send text directly to the PTY (for interactive input)
914 pub fn send_text(&mut self, text: &str) -> Result<()> {
915 debug!(bytes = text.len(), "PTY sending text");
916 self.send_bytes(text.as_bytes()).context("Failed to send text")?;
917 Ok(())
918 }
919
920 /// Send raw bytes directly to the PTY.
921 pub fn send_bytes(&mut self, bytes: &[u8]) -> Result<()> {
922 self.writer.write_all(bytes).context("Failed to send bytes")?;
923 self.writer.flush()?;
924 Ok(())
925 }
926
927 /// Send a special key sequence
928 pub fn send_special_key(&mut self, key: &str) -> Result<()> {
929 let bytes: &[u8] = match key {
930 "Enter" => b"\r",
931 "Tab" => b"\t",
932 "Backspace" => b"\x7F",
933 "Escape" => b"\x1B",
934 "Up" | "KeyUp" => b"\x1B[A",
935 "Down" | "KeyDown" => b"\x1B[B",
936 "Right" | "KeyRight" => b"\x1B[C",
937 "Left" | "KeyLeft" => b"\x1B[D",
938 "Home" => b"\x1B[H",
939 "End" => b"\x1B[F",
940 "PageUp" => b"\x1B[5~",
941 "PageDown" => b"\x1B[6~",
942 "Delete" => b"\x1B[3~",
943 "Insert" => b"\x1B[2~",
944 "CtrlC" | "Ctrl-C" => b"\x03",
945 "CtrlD" | "Ctrl-D" => b"\x04",
946 "CtrlZ" | "Ctrl-Z" => b"\x1A",
947 "CtrlL" | "Ctrl-L" => b"\x0C",
948 _ => return Err(anyhow!("Unknown special key: {key}")),
949 };
950
951 debug!("PTY sending special key: {} ({:?})", key, bytes);
952 self.send_bytes(bytes)?;
953 Ok(())
954 }
955
956 /// Snapshot the live terminal screen — a stable, consolidated view of what a
957 /// human would currently see, with cursor moves/redraws applied and ANSI
958 /// stripped. `max_lines` of 0 returns the full screen buffer. This is the
959 /// foundation for piloting interactive TUIs (the `claude` CLI, vim, ...).
960 pub fn live_snapshot(&self, max_lines: usize) -> Vec<String> {
961 // Recover from a poisoned lock instead of returning a blank screen
962 // forever — a single panic while holding `live` would otherwise wedge
963 // every future snapshot (and `wait_for_turn`/`screen` with it).
964 self.live.lock().unwrap_or_else(std::sync::PoisonError::into_inner).snapshot(max_lines)
965 }
966
967 /// Whether the live terminal is currently on the alternate screen buffer
968 /// (a full-screen app like vim/htop/less is running).
969 pub fn live_in_alt_screen(&self) -> bool {
970 self.live.lock().unwrap_or_else(std::sync::PoisonError::into_inner).in_alt_screen()
971 }
972
973 /// Cursor position `(row, col)` on the live screen (0-based), so a piloting
974 /// agent knows where focus is in a menu/form. `(0, 0)` if the lock is poisoned.
975 pub fn live_cursor_position(&self) -> (u16, u16) {
976 self.live.lock().unwrap_or_else(std::sync::PoisonError::into_inner).cursor_position()
977 }
978
979 /// Diff the live screen against what the client last saw, returning only the
980 /// changed lines when the change is small (huge token savings over many
981 /// polls). Updates the baseline as a side effect. See [`ScreenUpdate`].
982 pub fn live_snapshot_diff(&self, max_lines: usize, threshold: usize) -> ScreenUpdate {
983 self.live
984 .lock()
985 .unwrap_or_else(std::sync::PoisonError::into_inner)
986 .snapshot_diff(max_lines, threshold)
987 }
988
989 /// Resize the terminal
990 pub fn resize(&mut self, cols: u16, rows: u16) -> Result<()> {
991 debug!("PTY resizing to {}x{}", cols, rows);
992
993 let new_size = PtySize { rows, cols, pixel_width: 0, pixel_height: 0 };
994
995 self.master.resize(new_size).context("Failed to resize PTY")?;
996
997 // Keep the live emulator's viewport in lockstep with the PTY size.
998 self.live.lock().unwrap_or_else(std::sync::PoisonError::into_inner).resize(rows, cols);
999
1000 self.size = new_size;
1001 Ok(())
1002 }
1003
1004 /// Get current terminal size
1005 pub fn get_size(&self) -> (u16, u16) {
1006 (self.size.cols, self.size.rows)
1007 }
1008
1009 /// Check if the shell is still alive
1010 pub fn is_alive(&mut self) -> bool {
1011 self.child.try_wait().is_ok_and(|status| status.is_none())
1012 }
1013}
1014
1015/// Thread-safe wrapper for `PtyShell`
1016pub type SharedPtyShell = Arc<Mutex<Option<PtyShell>>>;
1017
1018/// Create a new shared PTY shell
1019pub fn create_shared_pty(initial_dir: &Path, restricted_mode: bool) -> Result<SharedPtyShell> {
1020 let shell = PtyShell::new(initial_dir, restricted_mode)?;
1021 Ok(Arc::new(Mutex::new(Some(shell))))
1022}
1023
1024/// Split a byte slice into the longest valid-UTF-8 prefix (decoded) and the
1025/// trailing bytes of an incomplete final char (to carry into the next read).
1026/// Genuinely invalid sequences become a single U+FFFD and are consumed, so the
1027/// carry only ever holds a recoverable incomplete tail and never grows unbounded.
1028///
1029/// Iterative on purpose: a recursive version blows the reader thread's stack on
1030/// a burst of invalid bytes (e.g. `cat`-ing a binary), which aborts under
1031/// `panic = "abort"`. This stays O(1) in stack depth regardless of input.
1032fn decode_keep_incomplete(bytes: &[u8]) -> (String, Vec<u8>) {
1033 let mut decoded = String::new();
1034 let mut rest = bytes;
1035 loop {
1036 match std::str::from_utf8(rest) {
1037 Ok(text) => {
1038 decoded.push_str(text);
1039 return (decoded, Vec::new());
1040 }
1041 Err(error) => {
1042 let valid = error.valid_up_to();
1043 // `rest[..valid]` is valid UTF-8 by construction (no allocation
1044 // for the borrowed case).
1045 decoded.push_str(&String::from_utf8_lossy(&rest[..valid]));
1046 match error.error_len() {
1047 // Incomplete final char: keep the tail for the next read.
1048 None => return (decoded, rest[valid..].to_vec()),
1049 // Genuinely invalid run: emit one replacement and skip past it.
1050 Some(bad) => {
1051 decoded.push('\u{FFFD}');
1052 rest = &rest[valid + bad..];
1053 }
1054 }
1055 }
1056 }
1057 }
1058}
1059
1060#[cfg(test)]
1061mod tests {
1062 use super::*;
1063 use tempfile::TempDir;
1064
1065 #[test]
1066 fn decode_passes_clean_utf8_through() {
1067 let (text, carry) = decode_keep_incomplete("hello ✓ café".as_bytes());
1068 assert_eq!(text, "hello ✓ café");
1069 assert!(carry.is_empty());
1070 }
1071
1072 #[test]
1073 fn decode_carries_incomplete_trailing_char_across_reads() {
1074 // "é" is 0xC3 0xA9; a read boundary splits it. The first half must NOT
1075 // turn into U+FFFD — it's held and completed on the next read.
1076 let (text1, carry1) = decode_keep_incomplete(&[b'a', 0xC3]);
1077 assert_eq!(text1, "a");
1078 assert_eq!(carry1, vec![0xC3]);
1079
1080 let mut next = carry1;
1081 next.push(0xA9); // continuation byte arrives in the next read
1082 let (text2, carry2) = decode_keep_incomplete(&next);
1083 assert_eq!(text2, "é");
1084 assert!(carry2.is_empty());
1085 }
1086
1087 #[test]
1088 fn decode_replaces_genuinely_invalid_bytes() {
1089 // 0xFF is never valid UTF-8 and is not an incomplete prefix.
1090 let (text, carry) = decode_keep_incomplete(&[b'x', 0xFF, b'y']);
1091 assert_eq!(text, "x\u{FFFD}y");
1092 assert!(carry.is_empty());
1093 }
1094
1095 #[test]
1096 fn decode_handles_large_invalid_burst_without_overflow() {
1097 // Regression: the first (recursive) version blew the stack — one frame
1098 // per invalid byte — and aborted under panic="abort" on a binary dump.
1099 // The iterative version must handle far more than the 4KB reader buffer.
1100 let big = vec![0xFFu8; 200_000];
1101 let (text, carry) = decode_keep_incomplete(&big);
1102 assert!(carry.is_empty());
1103 assert_eq!(text.chars().filter(|&c| c == '\u{FFFD}').count(), 200_000);
1104 }
1105
1106 #[test]
1107 fn prompt_detection_is_suffix_anchored() {
1108 // Per-shell nonce + exit code after the arrow: `──➤<nonce>:<code>`.
1109 let end = "──➤deadbeefcafe0001";
1110 // real prompt on the last line -> complete
1111 assert!(PtyShell::check_prompt_complete("out\nmore\n◉ /home/x──➤deadbeefcafe0001:0 ", end));
1112 // prompt with trailing ANSI erase -> still complete
1113 assert!(PtyShell::check_prompt_complete("◉ /home/x──➤deadbeefcafe0001:0 \u{1b}[K", end));
1114 // the bug: ◉ and ──➤ appear MID-output, last line is normal -> NOT complete
1115 assert!(!PtyShell::check_prompt_complete(
1116 "menu: ◉ start ──➤deadbeefcafe0001 stop\nstill running",
1117 end
1118 ));
1119 // command echoed after the prompt (not the waiting prompt) -> not complete
1120 assert!(!PtyShell::check_prompt_complete("◉ /home/x──➤deadbeefcafe0001:0 ls -la", end));
1121 // no prompt at all
1122 assert!(!PtyShell::check_prompt_complete("just some output\n", end));
1123 // the nonce's payoff: output that forges the bare glyphs but NOT the
1124 // nonce (or omits the exit code) must not be mistaken for the prompt.
1125 assert!(!PtyShell::check_prompt_complete("◉ /fake──➤ ", end));
1126 assert!(!PtyShell::check_prompt_complete("◉ /fake──➤wrongnonce:0 ", end));
1127 assert!(!PtyShell::check_prompt_complete("◉ /home/x──➤deadbeefcafe0001 ", end));
1128 }
1129
1130 #[test]
1131 fn parses_exit_code_from_prompt_marker() {
1132 let end = "──➤cafe1234";
1133 assert_eq!(PtyShell::parse_prompt_exit_code("ok\n◉ /x──➤cafe1234:0 ", end), Some(0));
1134 assert_eq!(PtyShell::parse_prompt_exit_code("oops\n◉ /x──➤cafe1234:1 ", end), Some(1));
1135 assert_eq!(PtyShell::parse_prompt_exit_code("◉ /x──➤cafe1234:42 ", end), Some(42));
1136 // still running (no completion prompt yet) -> None
1137 assert_eq!(PtyShell::parse_prompt_exit_code("building...", end), None);
1138 // forged with the wrong nonce -> None
1139 assert_eq!(PtyShell::parse_prompt_exit_code("◉ /x──➤wrong:0 ", end), None);
1140 }
1141
1142 use proptest::prelude::*;
1143
1144 proptest! {
1145 // Exit-code parsing runs on untrusted program output — arbitrary text
1146 // (and a forgeable marker) must never panic, only ever return None or a
1147 // parsed i32.
1148 #[test]
1149 fn parse_prompt_exit_code_never_panics(
1150 text in prop::collection::vec(any::<char>(), 0..256),
1151 nonce in "[0-9a-f]{0,16}",
1152 ) {
1153 let s: String = text.into_iter().collect();
1154 let end = format!("{WCGW_PROMPT_END}{nonce}");
1155 // Must not panic; any Option<i32> is acceptable.
1156 let _ = PtyShell::parse_prompt_exit_code(&s, &end);
1157 }
1158
1159 // A well-formed marker round-trips the embedded code for any i32.
1160 #[test]
1161 fn parse_prompt_exit_code_round_trips(code in any::<i32>(), nonce in "[0-9a-f]{1,16}") {
1162 let end = format!("{WCGW_PROMPT_END}{nonce}");
1163 let line = format!("output line\n{WCGW_PROMPT_PATTERN} /home/x{end}:{code} ");
1164 prop_assert_eq!(PtyShell::parse_prompt_exit_code(&line, &end), Some(code));
1165 }
1166 }
1167
1168 #[test]
1169 fn test_pty_shell_creation() -> Result<()> {
1170 let temp_dir = TempDir::new()?;
1171 let result = PtyShell::new(temp_dir.path(), false);
1172 assert!(result.is_ok(), "Failed to create PTY shell: {:?}", result.err());
1173 Ok(())
1174 }
1175
1176 #[test]
1177 fn test_pty_shell_echo() -> Result<()> {
1178 let temp_dir = TempDir::new()?;
1179 let mut shell = PtyShell::new(temp_dir.path(), false)?;
1180
1181 shell.send_command("echo 'hello pty'")?;
1182 let (output, _complete) = shell.read_output(2.0)?;
1183
1184 assert!(output.contains("hello pty"), "Output should contain 'hello pty': {output}");
1185 Ok(())
1186 }
1187
1188 #[test]
1189 fn completed_nonblocking_poll_is_level_triggered() -> Result<()> {
1190 let temp_dir = TempDir::new()?;
1191 let mut shell = PtyShell::new(temp_dir.path(), false)?;
1192
1193 shell.send_command("true")?;
1194 let deadline = Instant::now() + Duration::from_secs(2);
1195 while Instant::now() < deadline && !shell.poll_output_nonblocking() {
1196 std::thread::sleep(Duration::from_millis(10));
1197 }
1198
1199 assert!(!shell.command_running, "shell did not complete before the deadline");
1200 assert!(
1201 shell.poll_output_nonblocking(),
1202 "a later consumer must still observe the completed state"
1203 );
1204 Ok(())
1205 }
1206
1207 #[test]
1208 fn test_pty_shell_pwd() -> Result<()> {
1209 let temp_dir = TempDir::new()?;
1210 let mut shell = PtyShell::new(temp_dir.path(), false)?;
1211
1212 // Simply verify shell responds to pwd command
1213 // Use single quotes like echo test for consistency
1214 shell.send_command("pwd && echo 'pwd_done'")?;
1215 let (output, _complete) = shell.read_output(2.0)?;
1216
1217 // Verify the echo marker appears (proves command executed)
1218 assert!(output.contains("pwd_done"), "Output should contain 'pwd_done': {output}");
1219 Ok(())
1220 }
1221
1222 #[test]
1223 fn test_pty_resize() -> Result<()> {
1224 let temp_dir = TempDir::new()?;
1225 let mut shell = PtyShell::new(temp_dir.path(), false)?;
1226
1227 let result = shell.resize(120, 40);
1228 assert!(result.is_ok());
1229
1230 let (cols, rows) = shell.get_size();
1231 assert_eq!(cols, 120);
1232 assert_eq!(rows, 40);
1233 Ok(())
1234 }
1235
1236 /// `kill(pid, 0)` probes existence without sending a signal: 0 = the pid
1237 /// exists, -1 (ESRCH) = it's gone. A freshly-killed orphan lingers briefly
1238 /// as a zombie under init until reaped, so callers poll.
1239 #[cfg(unix)]
1240 fn pid_alive(pid: i32) -> bool {
1241 // SAFETY: kill(2) with signal 0 only checks permission/existence for an
1242 // integer pid; it sends nothing and touches no memory.
1243 unsafe { libc::kill(pid, 0) == 0 }
1244 }
1245
1246 #[cfg(unix)]
1247 #[test]
1248 fn dropping_shell_kills_orphaned_background_child() -> Result<()> {
1249 // A background job the shell starts must not outlive the shell. Spawn a
1250 // long `sleep` in the background, record its pid, drop the shell, and
1251 // confirm the sleep is gone - proving the group-kill reached past the
1252 // shell to a child `Child::kill` alone would have orphaned.
1253 let dir = TempDir::new()?;
1254 let pidfile = dir.path().join("child.pid");
1255 let mut shell = PtyShell::new(dir.path(), false)?;
1256
1257 shell
1258 .send_command(&format!("sleep 300 & echo $! > {} ; echo spawned", pidfile.display()))?;
1259
1260 // The PTY echoes the input line, so `out` containing "spawned" only proves
1261 // the line was sent - not that bash ran it. Under CI load bash can lag, so
1262 // poll the pidfile (it only appears once `echo $! > pidfile` executes) while
1263 // draining the PTY, instead of trusting the echo and racing the redirect.
1264 let mut child_pid = None;
1265 for _ in 0..150 {
1266 let _ = shell.read_output(0.1);
1267 if let Some(pid) =
1268 std::fs::read_to_string(&pidfile).ok().and_then(|c| c.trim().parse::<i32>().ok())
1269 {
1270 child_pid = Some(pid);
1271 break;
1272 }
1273 thread::sleep(Duration::from_millis(20));
1274 }
1275 let child_pid =
1276 child_pid.ok_or_else(|| anyhow!("background child pid never materialized"))?;
1277 assert!(pid_alive(child_pid), "background child should be running before drop");
1278
1279 drop(shell);
1280
1281 // The group-kill plus init reaping the orphan are async w.r.t. us; poll.
1282 let mut gone = false;
1283 for _ in 0..150 {
1284 if !pid_alive(child_pid) {
1285 gone = true;
1286 break;
1287 }
1288 thread::sleep(Duration::from_millis(20));
1289 }
1290 assert!(gone, "background child {child_pid} survived shell teardown (orphan leak)");
1291 Ok(())
1292 }
1293}