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