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