1use crate::detection::{Element, ElementDetector};
2use crate::pty::{key_to_escape_sequence, PtyError, PtyHandle};
3use crate::sync_utils::{mutex_lock_or_recover, rwlock_read_or_recover, rwlock_write_or_recover};
4use crate::terminal::{CursorPosition, VirtualTerminal};
5use chrono::{DateTime, Utc};
6use serde::{Deserialize, Serialize};
7use std::collections::{HashMap, VecDeque};
8use std::fs;
9use std::io::{BufReader, BufWriter};
10use std::path::PathBuf;
11use std::sync::{Arc, Mutex, RwLock};
12use std::time::Instant;
13use thiserror::Error;
14use uuid::Uuid;
15
16fn get_last_n<T: Clone>(queue: &VecDeque<T>, count: usize) -> Vec<T> {
17 let start = queue.len().saturating_sub(count);
18 queue.iter().skip(start).cloned().collect()
19}
20
21#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
22#[serde(transparent)]
23pub struct SessionId(String);
24
25impl SessionId {
26 pub fn new(id: impl Into<String>) -> Self {
27 Self(id.into())
28 }
29
30 pub fn generate() -> Self {
31 Self(Uuid::new_v4().to_string()[..8].to_string())
32 }
33
34 pub fn as_str(&self) -> &str {
35 &self.0
36 }
37}
38
39impl std::fmt::Display for SessionId {
40 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41 write!(f, "{}", self.0)
42 }
43}
44
45impl AsRef<str> for SessionId {
46 fn as_ref(&self) -> &str {
47 &self.0
48 }
49}
50
51impl From<String> for SessionId {
52 fn from(s: String) -> Self {
53 Self(s)
54 }
55}
56
57impl From<&str> for SessionId {
58 fn from(s: &str) -> Self {
59 Self(s.to_string())
60 }
61}
62
63#[derive(Clone, Debug)]
64pub struct RecordingFrame {
65 pub timestamp_ms: u64,
66 pub screen: String,
67}
68
69struct RecordingState {
70 is_recording: bool,
71 start_time: Instant,
72 frames: Vec<RecordingFrame>,
73}
74
75impl RecordingState {
76 fn new() -> Self {
77 Self {
78 is_recording: false,
79 start_time: Instant::now(),
80 frames: Vec::new(),
81 }
82 }
83}
84
85#[derive(Clone, Debug)]
86pub struct TraceEntry {
87 pub timestamp_ms: u64,
88 pub action: String,
89 pub details: Option<String>,
90}
91
92struct TraceState {
93 is_tracing: bool,
94 start_time: Instant,
95 entries: VecDeque<TraceEntry>,
96}
97
98impl TraceState {
99 fn new() -> Self {
100 Self {
101 is_tracing: false,
102 start_time: Instant::now(),
103 entries: VecDeque::new(),
104 }
105 }
106}
107
108pub struct RecordingStatus {
109 pub is_recording: bool,
110 pub frame_count: usize,
111 pub duration_ms: u64,
112}
113
114#[derive(Clone, Debug)]
115pub struct ErrorEntry {
116 pub timestamp: String,
117 pub message: String,
118 pub source: String,
119}
120
121struct ErrorState {
122 entries: VecDeque<ErrorEntry>,
123}
124
125impl ErrorState {
126 fn new() -> Self {
127 Self {
128 entries: VecDeque::new(),
129 }
130 }
131}
132
133#[derive(Error, Debug)]
134pub enum SessionError {
135 #[error("Session not found: {0}")]
136 NotFound(String),
137 #[error("No active session")]
138 NoActiveSession,
139 #[error("PTY error: {0}")]
140 Pty(#[from] PtyError),
141 #[error("Element not found: {0}")]
142 ElementNotFound(String),
143 #[error("Invalid key: {0}")]
144 InvalidKey(String),
145}
146
147#[derive(Clone, Copy)]
148enum ModifierKey {
149 Ctrl,
150 Alt,
151 Shift,
152 Meta,
153}
154
155impl ModifierKey {
156 fn from_str(key: &str) -> Option<Self> {
157 match key.to_lowercase().as_str() {
158 "ctrl" | "control" => Some(Self::Ctrl),
159 "alt" => Some(Self::Alt),
160 "shift" => Some(Self::Shift),
161 "meta" | "cmd" | "command" | "win" | "super" => Some(Self::Meta),
162 _ => None,
163 }
164 }
165}
166
167#[derive(Default)]
168struct ModifierState {
169 ctrl: bool,
170 alt: bool,
171 shift: bool,
172 meta: bool,
173}
174
175impl ModifierState {
176 fn set(&mut self, key: ModifierKey, value: bool) {
177 match key {
178 ModifierKey::Ctrl => self.ctrl = value,
179 ModifierKey::Alt => self.alt = value,
180 ModifierKey::Shift => self.shift = value,
181 ModifierKey::Meta => self.meta = value,
182 }
183 }
184}
185
186pub struct Session {
187 pub id: SessionId,
188 pub command: String,
189 pub created_at: DateTime<Utc>,
190 pty: PtyHandle,
191 terminal: VirtualTerminal,
192 detector: ElementDetector,
193 cached_elements: Vec<Element>,
194 recording: RecordingState,
195 trace: TraceState,
196 held_modifiers: ModifierState,
197 errors: ErrorState,
198}
199
200impl Session {
201 fn new(id: SessionId, command: String, pty: PtyHandle, cols: u16, rows: u16) -> Self {
202 Self {
203 id,
204 command,
205 created_at: Utc::now(),
206 pty,
207 terminal: VirtualTerminal::new(cols, rows),
208 detector: ElementDetector::new(),
209 cached_elements: Vec::new(),
210 recording: RecordingState::new(),
211 trace: TraceState::new(),
212 held_modifiers: ModifierState::default(),
213 errors: ErrorState::new(),
214 }
215 }
216
217 pub fn pid(&self) -> Option<u32> {
218 self.pty.pid()
219 }
220
221 pub fn is_running(&mut self) -> bool {
222 self.pty.is_running()
223 }
224
225 pub fn size(&self) -> (u16, u16) {
226 self.terminal.size()
227 }
228
229 pub fn update(&mut self) -> Result<(), SessionError> {
230 let mut buf = [0u8; 4096];
231
232 loop {
233 match self.pty.try_read(&mut buf, 10) {
234 Ok(0) => break,
235 Ok(n) => {
236 self.terminal.process(&buf[..n]);
237 }
238 Err(_) => break,
239 }
240 }
241
242 Ok(())
243 }
244
245 pub fn screen_text(&self) -> String {
246 self.terminal.screen_text()
247 }
248
249 pub fn cursor(&self) -> CursorPosition {
250 self.terminal.cursor()
251 }
252
253 pub fn detect_elements(&mut self) -> &[Element] {
254 let screen_text = self.terminal.screen_text();
255 let screen_buffer = self.terminal.screen_buffer();
256 self.cached_elements = self.detector.detect(&screen_text, Some(&screen_buffer));
257 &self.cached_elements
258 }
259
260 pub fn cached_elements(&self) -> &[Element] {
261 &self.cached_elements
262 }
263
264 pub fn find_element(&self, element_ref: &str) -> Option<&Element> {
265 self.detector
266 .find_by_ref(&self.cached_elements, element_ref)
267 }
268
269 pub fn keystroke(&self, key: &str) -> Result<(), SessionError> {
270 let seq =
271 key_to_escape_sequence(key).ok_or_else(|| SessionError::InvalidKey(key.to_string()))?;
272 self.pty.write(&seq)?;
273 Ok(())
274 }
275
276 pub fn keydown(&mut self, key: &str) -> Result<(), SessionError> {
277 let modifier = ModifierKey::from_str(key).ok_or_else(|| {
278 SessionError::InvalidKey(format!(
279 "{}. Only modifier keys (Ctrl, Alt, Shift, Meta) can be held",
280 key
281 ))
282 })?;
283 self.held_modifiers.set(modifier, true);
284 Ok(())
285 }
286
287 pub fn keyup(&mut self, key: &str) -> Result<(), SessionError> {
288 let modifier = ModifierKey::from_str(key).ok_or_else(|| {
289 SessionError::InvalidKey(format!(
290 "{}. Only modifier keys (Ctrl, Alt, Shift, Meta) can be released",
291 key
292 ))
293 })?;
294 self.held_modifiers.set(modifier, false);
295 Ok(())
296 }
297
298 pub fn type_text(&self, text: &str) -> Result<(), SessionError> {
299 self.pty.write_str(text)?;
300 Ok(())
301 }
302
303 pub fn click(&mut self, element_ref: &str) -> Result<(), SessionError> {
304 self.update()?;
305 self.detect_elements();
306
307 let element = self
308 .find_element(element_ref)
309 .ok_or_else(|| SessionError::ElementNotFound(element_ref.to_string()))?;
310
311 match element.element_type.as_str() {
312 "checkbox" | "radio" => {
313 self.pty.write(b" ")?;
314 }
315 _ => {
316 self.pty.write(b"\r")?;
317 }
318 }
319
320 Ok(())
321 }
322
323 pub fn fill(&mut self, element_ref: &str, value: &str) -> Result<(), SessionError> {
324 self.update()?;
325 self.detect_elements();
326
327 let _element = self
328 .find_element(element_ref)
329 .ok_or_else(|| SessionError::ElementNotFound(element_ref.to_string()))?;
330
331 self.pty.write_str(value)?;
332
333 Ok(())
334 }
335
336 pub fn resize(&mut self, cols: u16, rows: u16) -> Result<(), SessionError> {
337 self.pty.resize(cols, rows)?;
338 self.terminal.resize(cols, rows);
339 Ok(())
340 }
341
342 pub fn kill(&mut self) -> Result<(), SessionError> {
343 self.pty.kill()?;
344 Ok(())
345 }
346
347 pub fn pty_write(&self, data: &[u8]) -> Result<(), SessionError> {
348 self.pty.write(data)?;
349 Ok(())
350 }
351
352 pub fn pty_try_read(&self, buf: &mut [u8], timeout_ms: i32) -> Result<usize, SessionError> {
353 self.pty
354 .try_read(buf, timeout_ms)
355 .map_err(SessionError::Pty)
356 }
357
358 pub fn start_recording(&mut self) {
359 self.recording.is_recording = true;
360 self.recording.start_time = Instant::now();
361 self.recording.frames.clear();
362
363 let screen = self.terminal.screen_text();
364 self.recording.frames.push(RecordingFrame {
365 timestamp_ms: 0,
366 screen,
367 });
368 }
369
370 pub fn stop_recording(&mut self) -> Vec<RecordingFrame> {
371 self.recording.is_recording = false;
372 std::mem::take(&mut self.recording.frames)
373 }
374
375 pub fn recording_status(&self) -> RecordingStatus {
376 RecordingStatus {
377 is_recording: self.recording.is_recording,
378 frame_count: self.recording.frames.len(),
379 duration_ms: if self.recording.is_recording {
380 self.recording.start_time.elapsed().as_millis() as u64
381 } else {
382 0
383 },
384 }
385 }
386
387 pub fn start_trace(&mut self) {
388 self.trace.is_tracing = true;
389 self.trace.start_time = Instant::now();
390 self.trace.entries.clear();
391 }
392
393 pub fn stop_trace(&mut self) {
394 self.trace.is_tracing = false;
395 }
396
397 pub fn is_tracing(&self) -> bool {
398 self.trace.is_tracing
399 }
400
401 pub fn get_trace_entries(&self, count: usize) -> Vec<TraceEntry> {
402 get_last_n(&self.trace.entries, count)
403 }
404
405 pub fn get_errors(&self, count: usize) -> Vec<ErrorEntry> {
406 get_last_n(&self.errors.entries, count)
407 }
408
409 pub fn error_count(&self) -> usize {
410 self.errors.entries.len()
411 }
412
413 pub fn clear_errors(&mut self) {
414 self.errors.entries.clear();
415 }
416
417 pub fn clear_console(&mut self) {
418 self.terminal.clear();
419 }
420}
421
422pub struct SessionManager {
423 sessions: RwLock<HashMap<SessionId, Arc<Mutex<Session>>>>,
424 active_session: RwLock<Option<SessionId>>,
425 persistence: SessionPersistence,
426}
427
428impl Default for SessionManager {
429 fn default() -> Self {
430 Self::new()
431 }
432}
433
434impl SessionManager {
435 pub fn new() -> Self {
436 let persistence = SessionPersistence::new();
437 let _ = persistence.cleanup_stale_sessions();
438
439 Self {
440 sessions: RwLock::new(HashMap::new()),
441 active_session: RwLock::new(None),
442 persistence,
443 }
444 }
445
446 #[allow(clippy::too_many_arguments)]
447 pub fn spawn(
448 &self,
449 command: &str,
450 args: &[String],
451 cwd: Option<&str>,
452 env: Option<&HashMap<String, String>>,
453 session_id: Option<String>,
454 cols: u16,
455 rows: u16,
456 ) -> Result<(SessionId, u32), SessionError> {
457 let id = session_id
458 .map(SessionId::new)
459 .unwrap_or_else(SessionId::generate);
460
461 let pty = PtyHandle::spawn(command, args, cwd, env, cols, rows)?;
462 let pid = pty.pid().unwrap_or(0);
463
464 let session = Session::new(id.clone(), command.to_string(), pty, cols, rows);
465 let session = Arc::new(Mutex::new(session));
466
467 let created_at = Utc::now().to_rfc3339();
468 let persisted = PersistedSession {
469 id: id.clone(),
470 command: command.to_string(),
471 pid,
472 created_at,
473 cols,
474 rows,
475 };
476
477 {
478 let mut sessions = rwlock_write_or_recover(&self.sessions);
479 sessions.insert(id.clone(), session);
480 }
481
482 {
483 let mut active = rwlock_write_or_recover(&self.active_session);
484 *active = Some(id.clone());
485 }
486
487 if let Err(e) = self.persistence.add_session(persisted) {
488 eprintln!("Warning: Failed to persist session metadata: {}", e);
489 }
490
491 Ok((id, pid))
492 }
493
494 pub fn get(&self, session_id: &str) -> Result<Arc<Mutex<Session>>, SessionError> {
495 let sessions = rwlock_read_or_recover(&self.sessions);
496 let id = SessionId::new(session_id);
497 sessions
498 .get(&id)
499 .cloned()
500 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))
501 }
502
503 pub fn active(&self) -> Result<Arc<Mutex<Session>>, SessionError> {
504 let active_id = {
505 let active = rwlock_read_or_recover(&self.active_session);
506 active.clone()
507 };
508
509 match active_id {
510 Some(id) => self.get(id.as_str()),
511 None => Err(SessionError::NoActiveSession),
512 }
513 }
514
515 pub fn resolve(&self, session_id: Option<&str>) -> Result<Arc<Mutex<Session>>, SessionError> {
516 match session_id {
517 Some(id) => self.get(id),
518 None => self.active(),
519 }
520 }
521
522 pub fn set_active(&self, session_id: &str) -> Result<(), SessionError> {
523 let _ = self.get(session_id)?;
524
525 let mut active = rwlock_write_or_recover(&self.active_session);
526 *active = Some(SessionId::new(session_id));
527 Ok(())
528 }
529
530 pub fn list(&self) -> Vec<SessionInfo> {
531 let session_refs: Vec<(SessionId, Arc<Mutex<Session>>)> = {
532 let sessions = rwlock_read_or_recover(&self.sessions);
533 sessions
534 .iter()
535 .map(|(id, session)| (id.clone(), Arc::clone(session)))
536 .collect()
537 };
538
539 session_refs
540 .into_iter()
541 .map(|(id, session)| match session.try_lock() {
542 Ok(mut sess) => SessionInfo {
543 id: id.clone(),
544 command: sess.command.clone(),
545 pid: sess.pid().unwrap_or(0),
546 running: sess.is_running(),
547 created_at: sess.created_at.to_rfc3339(),
548 size: sess.size(),
549 },
550 Err(_) => SessionInfo {
551 id: id.clone(),
552 command: "(busy)".to_string(),
553 pid: 0,
554 running: true,
555 created_at: "".to_string(),
556 size: (80, 24),
557 },
558 })
559 .collect()
560 }
561
562 pub fn kill(&self, session_id: &str) -> Result<(), SessionError> {
563 let id = SessionId::new(session_id);
564
565 let session = {
566 let mut sessions = rwlock_write_or_recover(&self.sessions);
567 let mut active = rwlock_write_or_recover(&self.active_session);
568
569 let session = sessions
570 .remove(&id)
571 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
572
573 if active.as_ref() == Some(&id) {
574 *active = None;
575 }
576
577 session
578 };
579
580 {
581 let mut sess = mutex_lock_or_recover(&session);
582 sess.kill()?;
583 }
584
585 if let Err(e) = self.persistence.remove_session(session_id) {
586 eprintln!("Warning: Failed to remove session from persistence: {}", e);
587 }
588
589 Ok(())
590 }
591
592 pub fn session_count(&self) -> usize {
593 rwlock_read_or_recover(&self.sessions).len()
594 }
595
596 pub fn active_session_id(&self) -> Option<SessionId> {
597 rwlock_read_or_recover(&self.active_session).clone()
598 }
599}
600
601#[derive(Debug, Clone)]
602pub struct SessionInfo {
603 pub id: SessionId,
604 pub command: String,
605 pub pid: u32,
606 pub running: bool,
607 pub created_at: String,
608 pub size: (u16, u16),
609}
610
611#[derive(Debug, Clone, Serialize, Deserialize)]
612pub struct PersistedSession {
613 pub id: SessionId,
614 pub command: String,
615 pub pid: u32,
616 pub created_at: String,
617 pub cols: u16,
618 pub rows: u16,
619}
620
621pub struct SessionPersistence {
622 path: PathBuf,
623}
624
625impl SessionPersistence {
626 pub fn new() -> Self {
627 let path = Self::sessions_file_path();
628 Self { path }
629 }
630
631 fn sessions_file_path() -> PathBuf {
632 let home = std::env::var("HOME")
633 .map(PathBuf::from)
634 .unwrap_or_else(|_| PathBuf::from("/tmp"));
635 let dir = home.join(".agent-tui");
636 dir.join("sessions.json")
637 }
638
639 fn ensure_dir(&self) -> std::io::Result<()> {
640 if let Some(parent) = self.path.parent() {
641 fs::create_dir_all(parent).map_err(|e| {
642 std::io::Error::new(
643 e.kind(),
644 format!("Failed to create directory '{}': {}", parent.display(), e),
645 )
646 })?;
647 }
648 Ok(())
649 }
650
651 pub fn load(&self) -> Vec<PersistedSession> {
652 if !self.path.exists() {
653 return Vec::new();
654 }
655
656 match fs::File::open(&self.path) {
657 Ok(file) => {
658 let reader = BufReader::new(file);
659 serde_json::from_reader(reader).unwrap_or_default()
660 }
661 Err(e) => {
662 eprintln!(
663 "Warning: Failed to open sessions file '{}': {}",
664 self.path.display(),
665 e
666 );
667 Vec::new()
668 }
669 }
670 }
671
672 pub fn save(&self, sessions: &[PersistedSession]) -> std::io::Result<()> {
673 self.ensure_dir()?;
674
675 let file = fs::File::create(&self.path).map_err(|e| {
676 std::io::Error::new(
677 e.kind(),
678 format!("Failed to create file '{}': {}", self.path.display(), e),
679 )
680 })?;
681 let writer = BufWriter::new(file);
682 serde_json::to_writer_pretty(writer, sessions).map_err(|e| {
683 std::io::Error::other(format!(
684 "Failed to write sessions to '{}': {}",
685 self.path.display(),
686 e
687 ))
688 })?;
689 Ok(())
690 }
691
692 pub fn add_session(&self, session: PersistedSession) -> std::io::Result<()> {
693 let mut sessions = self.load();
694
695 sessions.retain(|s| s.id != session.id);
696 sessions.push(session);
697
698 self.save(&sessions)
699 }
700
701 pub fn remove_session(&self, session_id: &str) -> std::io::Result<()> {
702 let mut sessions = self.load();
703 sessions.retain(|s| s.id.as_str() != session_id);
704 self.save(&sessions)
705 }
706
707 pub fn cleanup_stale_sessions(&self) -> std::io::Result<usize> {
708 let sessions = self.load();
709 let mut cleaned = 0;
710
711 let active_sessions: Vec<PersistedSession> = sessions
712 .into_iter()
713 .filter(|s| {
714 let running = is_process_running(s.pid);
715 if !running {
716 cleaned += 1;
717 }
718 running
719 })
720 .collect();
721
722 self.save(&active_sessions)?;
723 Ok(cleaned)
724 }
725}
726
727impl Default for SessionPersistence {
728 fn default() -> Self {
729 Self::new()
730 }
731}
732
733fn is_process_running(pid: u32) -> bool {
734 unsafe { libc::kill(pid as i32, 0) == 0 }
735}
736
737impl From<&SessionInfo> for PersistedSession {
738 fn from(info: &SessionInfo) -> Self {
739 PersistedSession {
740 id: info.id.clone(),
741 command: info.command.clone(),
742 pid: info.pid,
743 created_at: info.created_at.clone(),
744 cols: info.size.0,
745 rows: info.size.1,
746 }
747 }
748}
749
750#[cfg(test)]
751mod persistence_tests {
752 use super::*;
753
754 #[test]
755 fn test_persisted_session_serialization() {
756 let session = PersistedSession {
757 id: SessionId::new("test123"),
758 command: "bash".to_string(),
759 pid: 12345,
760 created_at: "2024-01-01T00:00:00Z".to_string(),
761 cols: 80,
762 rows: 24,
763 };
764
765 let json = serde_json::to_string(&session).unwrap();
766 let parsed: PersistedSession = serde_json::from_str(&json).unwrap();
767
768 assert_eq!(parsed.id, session.id);
769 assert_eq!(parsed.command, session.command);
770 assert_eq!(parsed.pid, session.pid);
771 }
772
773 #[test]
774 fn test_is_process_running() {
775 let current_pid = std::process::id();
776 assert!(is_process_running(current_pid));
777
778 assert!(!is_process_running(999999999));
779 }
780}