1use 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
25pub struct SessionManager {
31 sessions: Mutex<HashMap<String, ShellSession>>,
32 config: ShellConfig,
33 next_id: AtomicU32,
34}
35
36struct 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#[derive(Debug, Clone, PartialEq)]
48pub enum SessionStatus {
49 Active,
50 Exited(Option<i32>),
51 Closed,
52}
53
54pub 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#[derive(Debug)]
67pub struct SendResult {
68 pub output: String,
69 pub status: String,
71}
72
73pub struct ShellSessionInfo {
75 pub id: String,
76 pub status: SessionStatus,
77 pub created_at: Instant,
78 pub last_active: Instant,
79}
80
81fn normalize_output(raw: &str) -> String {
87 strip_ansi(raw).replace("\r\n", "\n").replace('\r', "")
88}
89
90fn 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 } Some('b') => {
123 chars.next();
124 result.push('\x08');
125 } Some('0') => {
127 chars.next();
128 result.push('\0');
129 } Some('e') => {
131 chars.next();
132 tracing::warn!("blocked \\e escape sequence (raw ESC) in shell input");
133 }
134 Some('x') => {
135 chars.next(); 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 tracing::warn!(
151 "blocked \\x1b escape sequence (raw ESC) in shell input"
152 );
153 } else if byte >= 0x80 {
154 tracing::warn!("blocked \\x{hex:} high byte (>= 0x80) in shell input");
156 } else {
157 result.push(byte as char);
159 }
160 } else {
161 result.push('\\');
163 result.push('x');
164 result.push_str(&hex);
165 }
166 }
167 _ => {
168 result.push(ch);
170 }
171 }
172 } else {
173 result.push(ch);
174 }
175 }
176 result
177}
178
179fn 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
189async 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 let override_detector;
203 let effective_detector = if let Some(ms) = timeout_override {
204 override_detector = ReadinessDetector::new(
205 ReadinessStrategy::Hybrid,
206 &[], ms,
208 ms.saturating_mul(10).max(10_000), );
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 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 if let Some(tx) = tx_delta {
231 tx.send(normalize_output(&text));
232 }
233 }
234
235 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 if !pty.is_alive() {
247 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 if let Some(tx) = tx_delta {
256 tx.send(normalize_output(&remaining_text));
257 }
258 }
259 return (
261 normalize_output(&output),
262 status_string(&SessionStatus::Exited(None)),
263 );
264 }
265
266 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
279impl SessionManager {
284 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 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 {
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 let seq = self.next_id.fetch_add(1, Ordering::SeqCst) + 1;
326 let id = format!("shell_{:02}", seq);
327
328 let command = opts.command.unwrap_or_else(|| {
330 std::env::var("SHELL").unwrap_or_else(|_| {
331 if cfg!(windows) {
332 "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 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 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 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 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 {
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 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 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 if session.status != SessionStatus::Active {
448 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 let processed = process_input_escapes(input);
462 session.pty.write(processed.as_bytes())?;
463
464 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 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 {
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 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 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 drop(session);
525
526 Ok(final_output)
527 }
528
529 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 }
556
557 ids_to_reap
558 }
559
560 pub fn shutdown_all(&self) {
562 match self.sessions.lock() {
563 Ok(mut sessions) => {
564 sessions.drain();
565 }
567 Err(e) => {
568 tracing::error!("session lock poisoned: {e}");
569 }
570 }
571 }
572
573 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 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
604impl Drop for SessionManager {
609 fn drop(&mut self) {
610 self.shutdown_all();
611 }
612}
613
614pub 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#[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 #[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 #[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 let _ = mgr.close_session(&id).await;
716 }
717
718 #[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 #[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 let _ = mgr.close_session(&id1).await;
766 let _ = mgr.close_session(&id2).await;
767 }
768
769 #[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 #[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 #[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 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 assert_eq!(process_input_escapes(r"\q"), "\\q");
839 }
840
841 #[test]
842 fn test_escape_hex_partial() {
843 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 assert_eq!(process_input_escapes(r"\e"), "");
866 }
867
868 #[test]
869 fn test_escape_hex_1b_blocked() {
870 assert_eq!(process_input_escapes(r"\x1b"), "");
872 }
873
874 #[test]
875 fn test_escape_hex_high_byte_blocked() {
876 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 assert_eq!(process_input_escapes(r"\x7f"), "\x7f");
885 }
886}