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