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 crate::tools::utf8_stream::Utf8StreamDecoder;
24use super::readiness::{ReadinessDetector, ReadinessResult, ReadinessStrategy};
25
26pub struct SessionManager {
32 sessions: Mutex<HashMap<String, ShellSession>>,
33 config: ShellConfig,
34 next_id: AtomicU32,
35}
36
37struct ShellSession {
39 pty: PtyHandle,
40 decoder: Utf8StreamDecoder,
43 detector: ReadinessDetector,
44 created_at: Instant,
45 last_active: Instant,
46 idle_timeout: Duration,
47 status: SessionStatus,
48}
49
50#[derive(Debug, Clone, PartialEq)]
52pub enum SessionStatus {
53 Active,
54 Exited(Option<i32>),
55 Closed,
56}
57
58pub 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#[derive(Debug)]
71pub struct SendResult {
72 pub output: String,
73 pub status: String,
75}
76
77pub struct ShellSessionInfo {
79 pub id: String,
80 pub status: SessionStatus,
81 pub created_at: Instant,
82 pub last_active: Instant,
83}
84
85fn normalize_output(raw: &str) -> String {
91 strip_ansi(raw).replace("\r\n", "\n").replace('\r', "")
92}
93
94fn 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 } Some('b') => {
127 chars.next();
128 result.push('\x08');
129 } Some('0') => {
131 chars.next();
132 result.push('\0');
133 } Some('e') => {
135 chars.next();
136 tracing::warn!("blocked \\e escape sequence (raw ESC) in shell input");
137 }
138 Some('x') => {
139 chars.next(); 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 tracing::warn!(
155 "blocked \\x1b escape sequence (raw ESC) in shell input"
156 );
157 } else if byte >= 0x80 {
158 tracing::warn!("blocked \\x{hex:} high byte (>= 0x80) in shell input");
160 } else {
161 result.push(byte as char);
163 }
164 } else {
165 result.push('\\');
167 result.push('x');
168 result.push_str(&hex);
169 }
170 }
171 _ => {
172 result.push(ch);
174 }
175 }
176 } else {
177 result.push(ch);
178 }
179 }
180 result
181}
182
183fn 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
193async 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 let override_detector;
208 let effective_detector = if let Some(ms) = timeout_override {
209 override_detector = ReadinessDetector::new(
210 ReadinessStrategy::Hybrid,
211 &[], ms,
213 ms.saturating_mul(10).max(10_000), );
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 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 if let Some(tx) = tx_delta {
236 tx.send(normalize_output(&text));
237 }
238 }
239
240 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 if !pty.is_alive() {
252 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 if let Some(tx) = tx_delta {
262 tx.send(normalize_output(&remaining_text));
263 }
264 }
265 return (
267 normalize_output(&output),
268 status_string(&SessionStatus::Exited(None)),
269 );
270 }
271
272 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
285impl SessionManager {
290 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 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 {
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 let seq = self.next_id.fetch_add(1, Ordering::SeqCst) + 1;
332 let id = format!("shell_{:02}", seq);
333
334 let command = opts.command.unwrap_or_else(|| {
336 std::env::var("SHELL").unwrap_or_else(|_| {
337 if cfg!(windows) {
338 "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 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 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 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 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 {
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 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 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 if session.status != SessionStatus::Active {
457 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 let processed = process_input_escapes(input);
471 session.pty.write(processed.as_bytes())?;
472
473 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 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 {
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 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 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 drop(session);
537
538 Ok(final_output)
539 }
540
541 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 }
568
569 ids_to_reap
570 }
571
572 pub fn shutdown_all(&self) {
574 match self.sessions.lock() {
575 Ok(mut sessions) => {
576 sessions.drain();
577 }
579 Err(e) => {
580 tracing::error!("session lock poisoned: {e}");
581 }
582 }
583 }
584
585 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 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
616impl Drop for SessionManager {
621 fn drop(&mut self) {
622 self.shutdown_all();
623 }
624}
625
626pub 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#[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 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 #[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 #[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 let _ = mgr.close_session(&id).await;
730 }
731
732 #[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 #[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 #[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 let _ = mgr.close_session(&id1).await;
811 let _ = mgr.close_session(&id2).await;
812 }
813
814 #[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 #[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 #[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 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 assert_eq!(process_input_escapes(r"\q"), "\\q");
884 }
885
886 #[test]
887 fn test_escape_hex_partial() {
888 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 assert_eq!(process_input_escapes(r"\e"), "");
911 }
912
913 #[test]
914 fn test_escape_hex_1b_blocked() {
915 assert_eq!(process_input_escapes(r"\x1b"), "");
917 }
918
919 #[test]
920 fn test_escape_hex_high_byte_blocked() {
921 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 assert_eq!(process_input_escapes(r"\x7f"), "\x7f");
930 }
931}