Skip to main content

agent_engine/tools/shell/
session.rs

1//! Session manager — owns and manages active shell sessions.
2//!
3//! `SessionManager` is the core lifecycle engine for interactive PTY sessions.
4//! It provides thread-safe creation, I/O, and cleanup of `ShellSession`s, each
5//! backed by a `PtyHandle` and a `ReadinessDetector`.
6//!
7//! Design invariants:
8//! - The `Mutex` is held **only** for HashMap insert/remove — never during I/O.
9//! - Sessions are removed from the map before I/O and reinserted after, so the
10//!   lock is never contended by blocking reads/writes.
11//! - ANSI escape sequences are stripped from all output before returning.
12
13use std::collections::HashMap;
14use std::sync::atomic::{AtomicU32, Ordering};
15use std::sync::{Arc, Mutex};
16use std::time::{Duration, Instant};
17
18use crate::tools::strip_ansi;
19use crate::{Result, RuntimeError};
20
21use super::config::ShellConfig;
22use super::pty::PtyHandle;
23use super::readiness::{ReadinessDetector, ReadinessResult, ReadinessStrategy};
24
25// ---------------------------------------------------------------------------
26// Public types
27// ---------------------------------------------------------------------------
28
29/// Thread-safe manager for all active shell sessions.
30pub struct SessionManager {
31    sessions: Mutex<HashMap<String, ShellSession>>,
32    config: ShellConfig,
33    next_id: AtomicU32,
34}
35
36/// A single shell session — PTY handle, readiness detector, and metadata.
37struct ShellSession {
38    pty: PtyHandle,
39    detector: ReadinessDetector,
40    created_at: Instant,
41    last_active: Instant,
42    idle_timeout: Duration,
43    status: SessionStatus,
44}
45
46/// Lifecycle status of a session.
47#[derive(Debug, Clone, PartialEq)]
48pub enum SessionStatus {
49    Active,
50    Exited(Option<i32>),
51    Closed,
52}
53
54/// Options for creating a new session.
55pub struct SessionOpts {
56    pub command: Option<String>,
57    pub working_directory: Option<String>,
58    pub env: HashMap<String, String>,
59    pub rows: Option<u16>,
60    pub cols: Option<u16>,
61    pub readiness_timeout_ms: Option<u64>,
62    pub idle_timeout: Option<u64>,
63}
64
65/// Result of sending input to a session.
66#[derive(Debug)]
67pub struct SendResult {
68    pub output: String,
69    /// One of: `"active"`, `"exited"`, `"exited(N)"`, `"timeout"`
70    pub status: String,
71}
72
73/// Snapshot of session metadata (no mutable borrows needed).
74pub struct ShellSessionInfo {
75    pub id: String,
76    pub status: SessionStatus,
77    pub created_at: Instant,
78    pub last_active: Instant,
79}
80
81// ---------------------------------------------------------------------------
82// Core readiness polling loop
83// ---------------------------------------------------------------------------
84
85/// Normalize PTY output: strip ANSI escapes and convert \r\n → \n.
86fn normalize_output(raw: &str) -> String {
87    strip_ansi(raw).replace("\r\n", "\n").replace('\r', "")
88}
89
90/// Process escape sequences in input strings from the model.
91///
92/// The model sometimes sends literal two-character sequences like `\n` instead
93/// of actual control characters. This function converts common literal escapes
94/// to their real byte values as a defense-in-depth measure.
95fn process_input_escapes(input: &str) -> String {
96    let mut result = String::with_capacity(input.len());
97    let mut chars = input.chars().peekable();
98
99    while let Some(ch) = chars.next() {
100        if ch == '\\' {
101            match chars.peek() {
102                Some('n') => {
103                    chars.next();
104                    result.push('\n');
105                }
106                Some('r') => {
107                    chars.next();
108                    result.push('\r');
109                }
110                Some('t') => {
111                    chars.next();
112                    result.push('\t');
113                }
114                Some('\\') => {
115                    chars.next();
116                    result.push('\\');
117                }
118                Some('a') => {
119                    chars.next();
120                    result.push('\x07');
121                } // bell
122                Some('b') => {
123                    chars.next();
124                    result.push('\x08');
125                } // backspace
126                Some('0') => {
127                    chars.next();
128                    result.push('\0');
129                } // null
130                Some('e') => {
131                    chars.next();
132                    tracing::warn!("blocked \\e escape sequence (raw ESC) in shell input");
133                }
134                Some('x') => {
135                    chars.next(); // consume 'x'
136                    let mut hex = String::new();
137                    for _ in 0..2 {
138                        if let Some(&c) = chars.peek() {
139                            if c.is_ascii_hexdigit() {
140                                hex.push(c);
141                                chars.next();
142                            } else {
143                                break;
144                            }
145                        }
146                    }
147                    if let Ok(byte) = u8::from_str_radix(&hex, 16) {
148                        if byte == 0x1b {
149                            // Block ESC (ANSI escape initiator)
150                            tracing::warn!(
151                                "blocked \\x1b escape sequence (raw ESC) in shell input"
152                            );
153                        } else if byte >= 0x80 {
154                            // Block high bytes
155                            tracing::warn!("blocked \\x{hex:} high byte (>= 0x80) in shell input");
156                        } else {
157                            // Allow: 0x00-0x1a, 0x1c-0x1f (control chars except ESC), 0x20-0x7f
158                            result.push(byte as char);
159                        }
160                    } else {
161                        // Failed to parse — emit the original characters
162                        result.push('\\');
163                        result.push('x');
164                        result.push_str(&hex);
165                    }
166                }
167                _ => {
168                    // Unknown escape — pass through literally
169                    result.push(ch);
170                }
171            }
172        } else {
173            result.push(ch);
174        }
175    }
176    result
177}
178
179/// Format a `SessionStatus` as a status string.
180fn status_string(status: &SessionStatus) -> String {
181    match status {
182        SessionStatus::Active => "active".into(),
183        SessionStatus::Exited(Some(code)) => format!("exited({code})"),
184        SessionStatus::Exited(None) => "exited".into(),
185        SessionStatus::Closed => "closed".into(),
186    }
187}
188
189/// Wait for output from the PTY until the readiness detector signals completion.
190///
191/// Returns `(normalized_output, status_string)` where status is one of
192/// `"active"`, `"exited"`, `"exited(N)"`, or `"timeout"`.
193async fn wait_for_output(
194    pty: &mut PtyHandle,
195    detector: &ReadinessDetector,
196    timeout_override: Option<u64>,
197    tx_delta: Option<&crate::tools::output::DeltaSender>,
198    max_output: usize,
199) -> (String, String) {
200    // If a timeout override is provided, build a temporary detector with that
201    // silence timeout instead of using the session's detector.
202    let override_detector;
203    let effective_detector = if let Some(ms) = timeout_override {
204        override_detector = ReadinessDetector::new(
205            ReadinessStrategy::Hybrid,
206            &[], // no prompt patterns — falls back to Timeout strategy
207            ms,
208            ms.saturating_mul(10).max(10_000), // reasonable max timeout
209        );
210        &override_detector
211    } else {
212        detector
213    };
214
215    let mut output = String::new();
216    let start = Instant::now();
217    let mut last_output_time = Instant::now();
218    let poll_interval = Duration::from_millis(50);
219
220    loop {
221        // Try reading from PTY (async — waits up to poll_interval for data).
222        let bytes = pty.try_read_output(poll_interval).await;
223
224        if !bytes.is_empty() {
225            let text = String::from_utf8_lossy(&bytes);
226            output.push_str(&text);
227            last_output_time = Instant::now();
228
229            // Stream to TUI if requested (normalized)
230            if let Some(tx) = tx_delta {
231                tx.send(normalize_output(&text));
232            }
233        }
234
235        // Check if output exceeds the max size — truncate and return early.
236        if output.len() > max_output {
237            let mut trunc = max_output;
238            while trunc > 0 && !output.is_char_boundary(trunc) {
239                trunc -= 1;
240            }
241            output.truncate(trunc);
242            return (normalize_output(&output), "active".into());
243        }
244
245        // Check if process exited.
246        if !pty.is_alive() {
247            // Drain any remaining buffered output.
248            tokio::time::sleep(Duration::from_millis(50)).await;
249            let remaining = pty.try_read_output(Duration::from_millis(100)).await;
250            if !remaining.is_empty() {
251                let remaining_text = String::from_utf8_lossy(&remaining);
252                output.push_str(&remaining_text);
253
254                // Stream remaining output to TUI
255                if let Some(tx) = tx_delta {
256                    tx.send(normalize_output(&remaining_text));
257                }
258            }
259            // PtyHandle doesn't expose exit codes currently — use None
260            return (
261                normalize_output(&output),
262                status_string(&SessionStatus::Exited(None)),
263            );
264        }
265
266        // Evaluate readiness.
267        let silence_elapsed = last_output_time.elapsed();
268        let total_elapsed = start.elapsed();
269
270        match effective_detector.check(&output, silence_elapsed, total_elapsed) {
271            ReadinessResult::Ready => return (normalize_output(&output), "active".into()),
272            ReadinessResult::SilenceTimeout => return (normalize_output(&output), "active".into()),
273            ReadinessResult::MaxTimeout => return (normalize_output(&output), "timeout".into()),
274            ReadinessResult::Waiting => continue,
275        }
276    }
277}
278
279// ---------------------------------------------------------------------------
280// SessionManager implementation
281// ---------------------------------------------------------------------------
282
283impl SessionManager {
284    /// Create a new session manager backed by the given configuration.
285    pub fn new(config: ShellConfig) -> Arc<Self> {
286        Arc::new(Self {
287            sessions: Mutex::new(HashMap::new()),
288            config,
289            next_id: AtomicU32::new(0),
290        })
291    }
292
293    /// Create a new interactive shell session.
294    ///
295    /// Returns `(session_id, initial_output, status)` on success.
296    pub async fn create_session(
297        &self,
298        opts: SessionOpts,
299        tx_delta: Option<&crate::tools::output::DeltaSender>,
300        max_output: usize,
301    ) -> Result<(String, String, String)> {
302        self.create_session_cancellable(opts, tx_delta, max_output, None).await
303    }
304
305    pub(crate) async fn create_session_cancellable(
306        &self, opts: SessionOpts,
307        tx_delta: Option<&crate::tools::output::DeltaSender>, max_output: usize,
308        cancel: Option<&crate::CancellationToken>,
309    ) -> Result<(String, String, String)> {
310        // --- Check session limit ---
311        {
312            let sessions = self
313                .sessions
314                .lock()
315                .map_err(|e| RuntimeError::Tool(format!("session lock poisoned: {e}")))?;
316            if sessions.len() >= self.config.max_sessions {
317                return Err(RuntimeError::Tool(format!(
318                    "maximum session limit reached ({})",
319                    self.config.max_sessions
320                )));
321            }
322        }
323
324        // --- Generate ID (shell_01, shell_02, …) ---
325        let seq = self.next_id.fetch_add(1, Ordering::SeqCst) + 1;
326        let id = format!("shell_{:02}", seq);
327
328        // --- Resolve parameters with config defaults ---
329        let command = opts.command.unwrap_or_else(|| {
330            std::env::var("SHELL").unwrap_or_else(|_| {
331                if cfg!(windows) {
332                    // No $SHELL convention on Windows; PowerShell is the
333                    // native interactive default and always present.
334                    "powershell".into()
335                } else {
336                    "bash".into()
337                }
338            })
339        });
340        let rows = opts.rows.unwrap_or(self.config.default_rows);
341        let cols = opts.cols.unwrap_or(self.config.default_cols);
342
343        let idle_timeout = opts
344            .idle_timeout
345            .map(Duration::from_secs)
346            .unwrap_or(self.config.idle_timeout);
347
348        // --- Spawn PTY ---
349        if cancel.is_some_and(|c| c.is_cancelled()) {
350            return Err(RuntimeError::Tool("Shell canceled before launch".into()));
351        }
352        let probe_cwd = opts.working_directory.clone();
353        let probe_env = opts.env;
354        // The worker prepares only the probe; it never spawns a PTY command.
355        let probe = tokio::task::spawn_blocking(move ||
356            PtyHandle::prepare_scope(probe_cwd.as_deref(), &probe_env));
357        let prepared = if let Some(cancel) = cancel {
358            tokio::select! {
359                biased;
360                _ = cancel.cancelled() => return Err(RuntimeError::Tool("Shell canceled before launch".into())),
361                result = probe => result,
362            }
363        } else { probe.await }.map_err(|e| RuntimeError::Tool(e.to_string()))?;
364        if cancel.is_some_and(|c| c.is_cancelled()) {
365            return Err(RuntimeError::Tool("Shell canceled before launch".into()));
366        }
367        let mut pty = PtyHandle::spawn_prepared(
368            &command, opts.working_directory.as_deref(), rows, cols, prepared,
369        )?;
370
371        // --- Build readiness detector (with per-session overrides) ---
372        let silence_ms = opts
373            .readiness_timeout_ms
374            .unwrap_or(self.config.readiness_timeout_ms);
375        let detector = ReadinessDetector::new(
376            super::readiness::ReadinessStrategy::Hybrid,
377            &self.config.prompt_patterns,
378            silence_ms,
379            self.config.max_readiness_timeout_ms,
380        );
381
382        // --- Wait for initial output ---
383        // Give the process a moment to start producing output before polling.
384        // Without this, the silence timeout can fire before the process has
385        // had time to print anything (e.g. Python startup, shell rc files).
386        tokio::time::sleep(Duration::from_millis(200)).await;
387        let (initial_output, status_str) = wait_for_output(
388            &mut pty,
389            &detector,
390            opts.readiness_timeout_ms,
391            tx_delta,
392            max_output,
393        )
394        .await;
395
396        let now = Instant::now();
397        let status = if status_str.starts_with("exited") {
398            SessionStatus::Exited(None)
399        } else {
400            SessionStatus::Active
401        };
402
403        let session = ShellSession {
404            pty,
405            detector,
406            created_at: now,
407            last_active: now,
408            idle_timeout,
409            status,
410        };
411
412        // --- Insert into map ---
413        {
414            let mut sessions = self
415                .sessions
416                .lock()
417                .map_err(|e| RuntimeError::Tool(format!("session lock poisoned: {e}")))?;
418            sessions.insert(id.clone(), session);
419        }
420
421        Ok((id, initial_output, status_str))
422    }
423
424    /// Send input to an active session and return the output produced.
425    pub async fn send_input(
426        &self,
427        id: &str,
428        input: &str,
429        timeout_ms: Option<u64>,
430        tx_delta: Option<&crate::tools::output::DeltaSender>,
431        max_output: usize,
432    ) -> Result<SendResult> {
433        // --- Remove session from map (release lock before I/O) ---
434        let mut session = {
435            let mut sessions = self
436                .sessions
437                .lock()
438                .map_err(|e| RuntimeError::Tool(format!("session lock poisoned: {e}")))?;
439            sessions.remove(id).ok_or_else(|| {
440                RuntimeError::Tool(format!(
441                    "session {id} not found — it may have been closed, reaped, or is currently in use by another call"
442                ))
443            })?
444        };
445
446        // --- Reject if not active ---
447        if session.status != SessionStatus::Active {
448            // Reinsert so it can still be closed/inspected.
449            let s_str = status_string(&session.status);
450            let mut sessions = self
451                .sessions
452                .lock()
453                .map_err(|e| RuntimeError::Tool(format!("session lock poisoned: {e}")))?;
454            sessions.insert(id.to_string(), session);
455            return Err(RuntimeError::Tool(format!(
456                "session {id} is not active (status: {s_str})"
457            )));
458        }
459
460        // --- Write input (with escape sequence processing) ---
461        let processed = process_input_escapes(input);
462        session.pty.write(processed.as_bytes())?;
463
464        // --- Wait for output ---
465        let (output, status_str) = wait_for_output(
466            &mut session.pty,
467            &session.detector,
468            timeout_ms,
469            tx_delta,
470            max_output,
471        )
472        .await;
473
474        // --- Update metadata ---
475        session.last_active = Instant::now();
476        if !session.pty.is_alive() {
477            session.status = SessionStatus::Exited(None);
478        }
479
480        let result = SendResult {
481            output,
482            status: status_str,
483        };
484
485        // --- Reinsert into map ---
486        {
487            let mut sessions = self
488                .sessions
489                .lock()
490                .map_err(|e| RuntimeError::Tool(format!("session lock poisoned: {e}")))?;
491            sessions.insert(id.to_string(), session);
492        }
493
494        Ok(result)
495    }
496
497    /// Close a session, returning any final output.
498    ///
499    /// Idempotent — closing a non-existent session returns `Ok("")`.
500    pub async fn close_session(&self, id: &str) -> Result<String> {
501        let mut session = {
502            let mut sessions = self
503                .sessions
504                .lock()
505                .map_err(|e| RuntimeError::Tool(format!("session lock poisoned: {e}")))?;
506            match sessions.remove(id) {
507                Some(s) => s,
508                None => return Ok(String::new()),
509            }
510        };
511
512        // Read remaining output with a short timeout.
513        let remaining = session
514            .pty
515            .try_read_output(Duration::from_millis(100))
516            .await;
517        let final_output = if remaining.is_empty() {
518            String::new()
519        } else {
520            strip_ansi(&String::from_utf8_lossy(&remaining))
521        };
522
523        // PtyHandle::drop will kill the child process.
524        drop(session);
525
526        Ok(final_output)
527    }
528
529    /// Reap sessions that have been idle beyond their timeout.
530    ///
531    /// Returns the IDs of sessions that were removed.
532    pub fn reap_idle(&self) -> Vec<String> {
533        let mut sessions = match self.sessions.lock() {
534            Ok(s) => s,
535            Err(e) => {
536                tracing::error!("session lock poisoned: {e}");
537                return Vec::new();
538            }
539        };
540
541        let grace_period = Duration::from_secs(5);
542
543        let ids_to_reap: Vec<String> = sessions
544            .iter()
545            .filter(|(_, s)| {
546                let elapsed = s.last_active.elapsed();
547                elapsed > s.idle_timeout && elapsed > grace_period
548            })
549            .map(|(id, _)| id.clone())
550            .collect();
551
552        for id in &ids_to_reap {
553            sessions.remove(id);
554            // Dropped sessions clean up via PtyHandle::drop.
555        }
556
557        ids_to_reap
558    }
559
560    /// Shutdown all sessions immediately.
561    pub fn shutdown_all(&self) {
562        match self.sessions.lock() {
563            Ok(mut sessions) => {
564                sessions.drain();
565                // All PtyHandles dropped — children killed.
566            }
567            Err(e) => {
568                tracing::error!("session lock poisoned: {e}");
569            }
570        }
571    }
572
573    /// Number of sessions currently in the map.
574    pub fn active_count(&self) -> usize {
575        match self.sessions.lock() {
576            Ok(s) => s.len(),
577            Err(e) => {
578                tracing::error!("session lock poisoned: {e}");
579                0
580            }
581        }
582    }
583
584    /// Snapshot of all sessions.
585    pub fn list_sessions(&self) -> Vec<ShellSessionInfo> {
586        match self.sessions.lock() {
587            Ok(sessions) => sessions
588                .iter()
589                .map(|(id, s)| ShellSessionInfo {
590                    id: id.clone(),
591                    status: s.status.clone(),
592                    created_at: s.created_at,
593                    last_active: s.last_active,
594                })
595                .collect(),
596            Err(e) => {
597                tracing::error!("session lock poisoned: {e}");
598                Vec::new()
599            }
600        }
601    }
602}
603
604// ---------------------------------------------------------------------------
605// Drop implementation — cleanup on drop
606// ---------------------------------------------------------------------------
607
608impl Drop for SessionManager {
609    fn drop(&mut self) {
610        self.shutdown_all();
611    }
612}
613
614// ---------------------------------------------------------------------------
615// Background reaper
616// ---------------------------------------------------------------------------
617
618/// Spawn a background task that periodically reaps idle sessions.
619///
620/// The reaper runs every 30 seconds, checking for sessions whose last activity
621/// exceeds their configured idle timeout. Returns immediately — the task runs
622/// until the `CancellationToken` is canceled (or the process exits).
623pub fn start_reaper(
624    manager: Arc<SessionManager>,
625    cancel: tokio_util::sync::CancellationToken,
626) -> tokio::task::JoinHandle<()> {
627    tokio::spawn(async move {
628        let interval = Duration::from_secs(30);
629        loop {
630            tokio::select! {
631                _ = cancel.cancelled() => break,
632                _ = tokio::time::sleep(interval) => {
633                    let reaped = manager.reap_idle();
634                    for id in &reaped {
635                        tracing::info!(session_id = %id, "reaped idle shell session");
636                    }
637                }
638            }
639        }
640    })
641}
642
643// ---------------------------------------------------------------------------
644// Tests
645// ---------------------------------------------------------------------------
646
647#[cfg(test)]
648mod tests {
649    use super::*;
650
651    fn default_manager() -> Arc<SessionManager> {
652        SessionManager::new(ShellConfig::default())
653    }
654
655    fn opts_for(command: &str) -> SessionOpts {
656        SessionOpts {
657            command: Some(command.to_string()),
658            working_directory: None,
659            env: HashMap::new(),
660            rows: None,
661            cols: None,
662            readiness_timeout_ms: None,
663            idle_timeout: None,
664        }
665    }
666
667    #[tokio::test]
668    async fn canceled_creation_never_spawns_or_registers_a_shell() {
669        let mgr = default_manager();
670        let token = crate::CancellationToken::new();
671        token.cancel();
672        let error = mgr.create_session_cancellable(opts_for("echo should-not-run"), None,
673            256 * 1024, Some(&token)).await.unwrap_err().to_string();
674        assert!(error.contains("canceled before launch"), "{error}");
675        assert!(mgr.sessions.lock().unwrap().is_empty());
676    }
677
678    // 1. Create session with `echo hello` → output contains "hello"
679    #[tokio::test]
680    async fn test_create_session_echo_hello() {
681        let mgr = default_manager();
682        let (id, output, _status) = mgr
683            .create_session(opts_for("echo hello"), None, 256 * 1024)
684            .await
685            .expect("failed to create session");
686
687        assert!(id.starts_with("shell_"));
688        assert!(
689            output.contains("hello"),
690            "expected 'hello' in output, got: {output:?}"
691        );
692    }
693
694    // 2. Create bash session, send `echo test\n` → output contains "test"
695    #[tokio::test]
696    async fn test_send_input_echo() {
697        let mgr = default_manager();
698        let (id, _initial, _status) = mgr
699            .create_session(opts_for("bash"), None, 256 * 1024)
700            .await
701            .expect("failed to create session");
702
703        let result = mgr
704            .send_input(&id, "echo test\n", None, None, 256 * 1024)
705            .await
706            .expect("failed to send input");
707
708        assert!(
709            result.output.contains("test"),
710            "expected 'test' in output, got: {:?}",
711            result.output
712        );
713
714        // Clean up
715        let _ = mgr.close_session(&id).await;
716    }
717
718    // 3. Close session → idempotent (close twice is fine)
719    #[tokio::test]
720    async fn test_close_session_idempotent() {
721        let mgr = default_manager();
722        let (id, _, _status) = mgr
723            .create_session(opts_for("bash"), None, 256 * 1024)
724            .await
725            .expect("failed to create session");
726
727        let result1 = mgr.close_session(&id).await;
728        assert!(result1.is_ok(), "first close should succeed");
729
730        let result2 = mgr.close_session(&id).await;
731        assert!(
732            result2.is_ok(),
733            "second close should also succeed (idempotent)"
734        );
735        assert_eq!(result2.unwrap(), "", "second close returns empty string");
736    }
737
738    // 4. Max sessions limit → error on exceeding
739    #[tokio::test]
740    async fn test_max_sessions_limit() {
741        let config = ShellConfig {
742            max_sessions: 2,
743            ..Default::default()
744        };
745        let mgr = SessionManager::new(config);
746
747        let (id1, _, _s) = mgr
748            .create_session(opts_for("bash"), None, 256 * 1024)
749            .await
750            .expect("session 1");
751        let (id2, _, _s) = mgr
752            .create_session(opts_for("bash"), None, 256 * 1024)
753            .await
754            .expect("session 2");
755
756        let result = mgr.create_session(opts_for("bash"), None, 256 * 1024).await;
757        assert!(result.is_err(), "third session should fail");
758        let err_msg = format!("{}", result.unwrap_err());
759        assert!(
760            err_msg.contains("maximum session limit"),
761            "error should mention limit, got: {err_msg}"
762        );
763
764        // Clean up
765        let _ = mgr.close_session(&id1).await;
766        let _ = mgr.close_session(&id2).await;
767    }
768
769    // 5. Session not found → error
770    #[tokio::test]
771    async fn test_session_not_found() {
772        let mgr = default_manager();
773        let result = mgr
774            .send_input("shell_99", "hello\n", None, None, 256 * 1024)
775            .await;
776        assert!(result.is_err(), "send to non-existent session should fail");
777        let err_msg = format!("{}", result.unwrap_err());
778        assert!(
779            err_msg.contains("not found"),
780            "error should mention 'not found', got: {err_msg}"
781        );
782    }
783
784    // ── normalize_output tests ──
785
786    #[test]
787    fn test_normalize_output_crlf() {
788        assert_eq!(normalize_output("hello\r\nworld\r\n"), "hello\nworld\n");
789    }
790
791    #[test]
792    fn test_normalize_output_lone_cr() {
793        assert_eq!(normalize_output("abc\rdef"), "abcdef");
794    }
795
796    // ── process_input_escapes tests ──
797
798    #[test]
799    fn test_escape_newline() {
800        assert_eq!(process_input_escapes(r"hello\n"), "hello\n");
801    }
802
803    #[test]
804    fn test_escape_tab() {
805        assert_eq!(process_input_escapes(r"a\tb"), "a\tb");
806    }
807
808    #[test]
809    fn test_escape_ctrl_c() {
810        assert_eq!(process_input_escapes(r"\x03"), "\x03");
811    }
812
813    #[test]
814    fn test_escape_ctrl_d() {
815        assert_eq!(process_input_escapes(r"\x04"), "\x04");
816    }
817
818    #[test]
819    fn test_escape_literal_backslash() {
820        assert_eq!(process_input_escapes(r"a\\b"), "a\\b");
821    }
822
823    #[test]
824    fn test_escape_real_newline_passthrough() {
825        // If the model sends an actual newline (JSON parsed correctly), it passes through
826        assert_eq!(process_input_escapes("hello\n"), "hello\n");
827    }
828
829    #[test]
830    fn test_escape_mixed() {
831        assert_eq!(process_input_escapes(r"ls -la\n"), "ls -la\n");
832        assert_eq!(process_input_escapes(r"124\n"), "124\n");
833    }
834
835    #[test]
836    fn test_escape_unknown_sequence() {
837        // Unknown escapes pass through literally
838        assert_eq!(process_input_escapes(r"\q"), "\\q");
839    }
840
841    #[test]
842    fn test_escape_hex_partial() {
843        // Incomplete hex — pass through
844        assert_eq!(process_input_escapes(r"\xZZ"), "\\xZZ");
845    }
846
847    #[test]
848    fn test_escape_bell() {
849        assert_eq!(process_input_escapes(r"\a"), "\x07");
850    }
851
852    #[test]
853    fn test_escape_backspace() {
854        assert_eq!(process_input_escapes(r"\b"), "\x08");
855    }
856
857    #[test]
858    fn test_escape_null() {
859        assert_eq!(process_input_escapes(r"\0"), "\0");
860    }
861
862    #[test]
863    fn test_escape_esc_blocked() {
864        // \e should be blocked (produces empty string for that escape)
865        assert_eq!(process_input_escapes(r"\e"), "");
866    }
867
868    #[test]
869    fn test_escape_hex_1b_blocked() {
870        // \x1b should be blocked
871        assert_eq!(process_input_escapes(r"\x1b"), "");
872    }
873
874    #[test]
875    fn test_escape_hex_high_byte_blocked() {
876        // \x80 and above should be blocked
877        assert_eq!(process_input_escapes(r"\x80"), "");
878        assert_eq!(process_input_escapes(r"\xff"), "");
879    }
880
881    #[test]
882    fn test_escape_hex_del_allowed() {
883        // \x7f (DEL) should be allowed
884        assert_eq!(process_input_escapes(r"\x7f"), "\x7f");
885    }
886}