ai_session/core/
mod.rs

1//! Core session management functionality
2//!
3//! This module provides the foundational components for AI-optimized terminal session management.
4//! The core functionality includes session creation, lifecycle management, and AI context integration.
5//!
6//! # Key Features
7//!
8//! - **AISession**: Advanced terminal session with AI capabilities
9//! - **SessionManager**: Pool-based session management with automatic cleanup
10//! - **SessionConfig**: Comprehensive configuration for AI features and performance
11//! - **Context Integration**: Seamless integration with AI conversation context
12//!
13//! # Quick Start
14//!
15//! ```no_run
16//! use ai_session::{SessionManager, SessionConfig, ContextConfig};
17//! use tokio;
18//!
19//! #[tokio::main]
20//! async fn main() -> anyhow::Result<()> {
21//!     let manager = SessionManager::new();
22//!     
23//!     // Create a basic session
24//!     let session = manager.create_session().await?;
25//!     session.start().await?;
26//!     
27//!     // Send a command
28//!     session.send_input("echo 'Hello AI Session!'\n").await?;
29//!     
30//!     // Read the output
31//!     tokio::time::sleep(std::time::Duration::from_millis(300)).await;
32//!     let output = session.read_output().await?;
33//!     println!("Output: {}", String::from_utf8_lossy(&output));
34//!     
35//!     // Clean up
36//!     session.stop().await?;
37//!     Ok(())
38//! }
39//! ```
40//!
41//! # Advanced Configuration
42//!
43//! ```no_run
44//! use ai_session::{SessionManager, SessionConfig, ContextConfig};
45//! use std::collections::HashMap;
46//!
47//! #[tokio::main]
48//! async fn main() -> anyhow::Result<()> {
49//!     let manager = SessionManager::new();
50//!     
51//!     // Configure session with AI features
52//!     let mut config = SessionConfig::default();
53//!     config.enable_ai_features = true;
54//!     config.agent_role = Some("rust-developer".to_string());
55//!     config.context_config = ContextConfig {
56//!         max_tokens: 8192,
57//!         compression_threshold: 0.8,
58//!     };
59//!     
60//!     // Set environment variables
61//!     config.environment.insert("RUST_LOG".to_string(), "debug".to_string());
62//!     config.working_directory = "/path/to/project".into();
63//!     
64//!     let session = manager.create_session_with_config(config).await?;
65//!     session.start().await?;
66//!     
67//!     // Session is now ready for AI-enhanced development
68//!     Ok(())
69//! }
70//! ```
71
72use std::collections::HashMap;
73use std::path::PathBuf;
74use std::sync::Arc;
75use std::time::Duration;
76
77use anyhow::Result;
78use chrono::{DateTime, Utc};
79use dashmap::DashMap;
80use serde::{Deserialize, Serialize};
81use tokio::sync::RwLock;
82use uuid::Uuid;
83
84pub mod headless;
85pub mod lifecycle;
86pub mod process;
87pub mod pty;
88pub mod terminal;
89
90use crate::context::SessionContext;
91use crate::persistence::CommandRecord;
92
93/// Session error type
94#[derive(Debug, thiserror::Error)]
95pub enum SessionError {
96    #[error("Session not found: {0}")]
97    NotFound(SessionId),
98
99    #[error("Session already exists: {0}")]
100    AlreadyExists(SessionId),
101
102    #[error("PTY error: {0}")]
103    PtyError(String),
104
105    #[error("Process error: {0}")]
106    ProcessError(String),
107
108    #[error("IO error: {0}")]
109    IoError(#[from] std::io::Error),
110
111    #[error("Other error: {0}")]
112    Other(#[from] anyhow::Error),
113}
114
115/// Session result type
116pub type SessionResult<T> = std::result::Result<T, SessionError>;
117
118/// Unique session identifier
119#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)]
120pub struct SessionId(Uuid);
121
122impl Default for SessionId {
123    fn default() -> Self {
124        Self::new()
125    }
126}
127
128impl SessionId {
129    /// Create a new unique session ID
130    pub fn new() -> Self {
131        Self(Uuid::new_v4())
132    }
133
134    /// Create a new v4 UUID session ID
135    pub fn new_v4() -> Self {
136        Self(Uuid::new_v4())
137    }
138
139    /// Parse from string
140    pub fn parse_str(s: &str) -> Result<Self> {
141        Ok(Self(Uuid::parse_str(s)?))
142    }
143
144    /// Convert to string
145    pub fn to_string(&self) -> String {
146        self.0.to_string()
147    }
148
149    /// Get inner UUID
150    pub fn as_uuid(&self) -> &Uuid {
151        &self.0
152    }
153}
154
155impl std::fmt::Display for SessionId {
156    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
157        write!(f, "{}", self.0)
158    }
159}
160
161/// Session status
162#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
163pub enum SessionStatus {
164    /// Session is being initialized
165    #[default]
166    Initializing,
167    /// Session is running and ready
168    Running,
169    /// Session is paused
170    Paused,
171    /// Session is being terminated
172    Terminating,
173    /// Session has been terminated
174    Terminated,
175    /// Session encountered an error
176    Error,
177}
178
179/// Session configuration
180#[derive(Debug, Clone, Serialize, Deserialize)]
181#[serde(default)]
182pub struct SessionConfig {
183    /// Session name (optional)
184    pub name: Option<String>,
185    /// Working directory
186    pub working_directory: PathBuf,
187    /// Environment variables
188    pub environment: HashMap<String, String>,
189    /// Shell command to execute
190    pub shell: Option<String>,
191    /// Shell command to execute (alternative field for compatibility)
192    pub shell_command: Option<String>,
193    /// PTY size (rows, cols)
194    pub pty_size: (u16, u16),
195    /// Output buffer size in bytes
196    pub output_buffer_size: usize,
197    /// Session timeout (None for no timeout)
198    pub timeout: Option<Duration>,
199    /// Enable output compression
200    pub compress_output: bool,
201    /// Enable semantic output parsing
202    pub parse_output: bool,
203    /// Enable AI features
204    pub enable_ai_features: bool,
205    /// Context configuration
206    pub context_config: ContextConfig,
207    /// Agent role (optional)
208    pub agent_role: Option<String>,
209    /// Force headless (non-PTY) execution (useful for restricted sandboxes)
210    pub force_headless: bool,
211    /// Allow automatic fallback to headless mode when PTY creation fails
212    pub allow_headless_fallback: bool,
213}
214
215/// Context configuration for AI features
216#[derive(Debug, Clone, Serialize, Deserialize)]
217#[serde(default)]
218pub struct ContextConfig {
219    /// Maximum tokens for context
220    pub max_tokens: usize,
221    /// Compression threshold (0.0 to 1.0)
222    pub compression_threshold: f64,
223}
224
225impl Default for SessionConfig {
226    fn default() -> Self {
227        Self {
228            name: None,
229            working_directory: std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/")),
230            environment: HashMap::new(),
231            shell: None,
232            shell_command: None,
233            pty_size: (24, 80),
234            output_buffer_size: 1024 * 1024, // 1MB
235            timeout: None,
236            compress_output: true,
237            parse_output: true,
238            enable_ai_features: false,
239            context_config: ContextConfig::default(),
240            agent_role: None,
241            force_headless: false,
242            allow_headless_fallback: true,
243        }
244    }
245}
246
247impl Default for ContextConfig {
248    fn default() -> Self {
249        Self {
250            max_tokens: 4096,
251            compression_threshold: 0.8,
252        }
253    }
254}
255
256/// AI-optimized session
257pub struct AISession {
258    /// Unique session ID
259    pub id: SessionId,
260    /// Session configuration
261    pub config: SessionConfig,
262    /// Current status
263    pub status: RwLock<SessionStatus>,
264    /// Session context (AI state, history, etc.)
265    pub context: Arc<RwLock<SessionContext>>,
266    /// Process handle
267    process: Arc<RwLock<Option<process::ProcessHandle>>>,
268    /// Terminal handle (PTY or headless)
269    terminal: Arc<RwLock<Option<terminal::TerminalHandle>>>,
270    /// Creation time
271    pub created_at: DateTime<Utc>,
272    /// Last activity time
273    pub last_activity: Arc<RwLock<DateTime<Utc>>>,
274    /// Session metadata
275    pub metadata: Arc<RwLock<HashMap<String, serde_json::Value>>>,
276    /// Command history tracking
277    pub command_history: Arc<RwLock<Vec<CommandRecord>>>,
278    /// Command count
279    pub command_count: Arc<RwLock<usize>>,
280    /// Total tokens used
281    pub total_tokens: Arc<RwLock<usize>>,
282}
283
284impl AISession {
285    /// Create a new AI session
286    pub async fn new(config: SessionConfig) -> Result<Self> {
287        let id = SessionId::new();
288        let now = Utc::now();
289
290        Ok(Self {
291            id: id.clone(),
292            config,
293            status: RwLock::new(SessionStatus::Initializing),
294            context: Arc::new(RwLock::new(SessionContext::new(id))),
295            process: Arc::new(RwLock::new(None)),
296            terminal: Arc::new(RwLock::new(None)),
297            created_at: now,
298            last_activity: Arc::new(RwLock::new(now)),
299            metadata: Arc::new(RwLock::new(HashMap::new())),
300            command_history: Arc::new(RwLock::new(Vec::new())),
301            command_count: Arc::new(RwLock::new(0)),
302            total_tokens: Arc::new(RwLock::new(0)),
303        })
304    }
305
306    /// Create an AI session with a specific ID (for restoration)
307    pub async fn new_with_id(
308        id: SessionId,
309        config: SessionConfig,
310        created_at: DateTime<Utc>,
311    ) -> Result<Self> {
312        let now = Utc::now();
313
314        Ok(Self {
315            id: id.clone(),
316            config,
317            status: RwLock::new(SessionStatus::Initializing),
318            context: Arc::new(RwLock::new(SessionContext::new(id))),
319            process: Arc::new(RwLock::new(None)),
320            terminal: Arc::new(RwLock::new(None)),
321            created_at,
322            last_activity: Arc::new(RwLock::new(now)),
323            metadata: Arc::new(RwLock::new(HashMap::new())),
324            command_history: Arc::new(RwLock::new(Vec::new())),
325            command_count: Arc::new(RwLock::new(0)),
326            total_tokens: Arc::new(RwLock::new(0)),
327        })
328    }
329
330    /// Start the session
331    pub async fn start(&self) -> Result<()> {
332        lifecycle::start_session(self).await
333    }
334
335    /// Stop the session
336    pub async fn stop(&self) -> Result<()> {
337        lifecycle::stop_session(self).await
338    }
339
340    /// Send input to the session
341    pub async fn send_input(&self, input: &str) -> Result<()> {
342        let terminal_guard = self.terminal.read().await;
343        if let Some(terminal) = terminal_guard.as_ref() {
344            terminal.write(input.as_bytes()).await?;
345            *self.last_activity.write().await = Utc::now();
346            Ok(())
347        } else {
348            Err(anyhow::anyhow!("Session not started"))
349        }
350    }
351
352    /// Read output from the session
353    pub async fn read_output(&self) -> Result<Vec<u8>> {
354        let terminal = self.terminal.read().await;
355        if let Some(terminal) = terminal.as_ref() {
356            let output = terminal.read().await?;
357            *self.last_activity.write().await = Utc::now();
358            Ok(output)
359        } else {
360            Err(anyhow::anyhow!("Session not started"))
361        }
362    }
363
364    /// Get current session status
365    pub async fn status(&self) -> SessionStatus {
366        *self.status.read().await
367    }
368
369    /// Update session metadata
370    pub async fn set_metadata(&self, key: String, value: serde_json::Value) -> Result<()> {
371        self.metadata.write().await.insert(key, value);
372        Ok(())
373    }
374
375    /// Get session metadata
376    pub async fn get_metadata(&self, key: &str) -> Option<serde_json::Value> {
377        self.metadata.read().await.get(key).cloned()
378    }
379
380    /// Execute a command and record it in history
381    pub async fn execute_command(&self, command: &str) -> Result<String> {
382        let start_time = Utc::now();
383
384        // Send the command
385        self.send_input(&format!("{}\n", command)).await?;
386
387        // Wait for output
388        tokio::time::sleep(Duration::from_millis(500)).await;
389        let output_bytes = self.read_output().await?;
390        let output = String::from_utf8_lossy(&output_bytes).to_string();
391
392        // Record the command in history
393        let end_time = Utc::now();
394        let duration_ms = (end_time - start_time).num_milliseconds() as u64;
395
396        let record = CommandRecord {
397            command: command.to_string(),
398            timestamp: start_time,
399            exit_code: None, // TODO: Extract exit code from output
400            output_preview: if output.len() > 200 {
401                format!("{}...", &output[..200])
402            } else {
403                output.clone()
404            },
405            duration_ms,
406        };
407
408        // Update history and counters
409        self.command_history.write().await.push(record);
410        *self.command_count.write().await += 1;
411
412        Ok(output)
413    }
414
415    /// Add tokens to the session total
416    pub async fn add_tokens(&self, token_count: usize) {
417        *self.total_tokens.write().await += token_count;
418    }
419
420    /// Get command history
421    pub async fn get_command_history(&self) -> Vec<CommandRecord> {
422        self.command_history.read().await.clone()
423    }
424
425    /// Get command count
426    pub async fn get_command_count(&self) -> usize {
427        *self.command_count.read().await
428    }
429
430    /// Get total tokens used
431    pub async fn get_total_tokens(&self) -> usize {
432        *self.total_tokens.read().await
433    }
434
435    /// Clear command history (keep recent N commands)
436    pub async fn trim_command_history(&self, keep_recent: usize) {
437        let mut history = self.command_history.write().await;
438        if history.len() > keep_recent {
439            let start_index = history.len() - keep_recent;
440            history.drain(0..start_index);
441        }
442    }
443}
444
445/// AI-optimized session manager for creating and managing multiple terminal sessions.
446///
447/// The `SessionManager` provides a centralized way to create, track, and manage AI-enhanced
448/// terminal sessions. It includes automatic cleanup, session restoration, and efficient
449/// resource management.
450///
451/// # Features
452///
453/// - **Session Pooling**: Efficient management of multiple concurrent sessions
454/// - **Automatic Cleanup**: Garbage collection of terminated sessions
455/// - **Session Restoration**: Restore sessions from persistent storage
456/// - **Resource Management**: Automatic cleanup and memory management
457///
458/// # Examples
459///
460/// ## Basic Session Management
461///
462/// ```no_run
463/// use ai_session::{SessionManager, SessionConfig};
464///
465/// #[tokio::main]
466/// async fn main() -> anyhow::Result<()> {
467///     let manager = SessionManager::new();
468///     
469///     // Create multiple sessions
470///     let session1 = manager.create_session().await?;
471///     let session2 = manager.create_session().await?;
472///     
473///     session1.start().await?;
474///     session2.start().await?;
475///     
476///     // List all active sessions
477///     let session_ids = manager.list_sessions();
478///     println!("Active sessions: {}", session_ids.len());
479///     
480///     // Clean up
481///     manager.remove_session(&session1.id).await?;
482///     manager.remove_session(&session2.id).await?;
483///     
484///     Ok(())
485/// }
486/// ```
487///
488/// ## Custom Configuration
489///
490/// ```no_run
491/// use ai_session::{SessionManager, SessionConfig, ContextConfig};
492///
493/// #[tokio::main]
494/// async fn main() -> anyhow::Result<()> {
495///     let manager = SessionManager::new();
496///     
497///     // Configure for AI development agent
498///     let mut config = SessionConfig::default();
499///     config.enable_ai_features = true;
500///     config.agent_role = Some("backend-developer".to_string());
501///     config.working_directory = "/project/backend".into();
502///     config.context_config = ContextConfig {
503///         max_tokens: 8192,
504///         compression_threshold: 0.8,
505///     };
506///     
507///     let session = manager.create_session_with_config(config).await?;
508///     session.start().await?;
509///     
510///     // Session is optimized for AI backend development
511///     Ok(())
512/// }
513/// ```
514///
515/// ## Session Persistence
516///
517/// ```no_run
518/// use ai_session::{SessionManager, SessionConfig, SessionId};
519/// use chrono::Utc;
520///
521/// #[tokio::main]
522/// async fn main() -> anyhow::Result<()> {
523///     let manager = SessionManager::new();
524///     
525///     // Create session
526///     let config = SessionConfig::default();
527///     let session = manager.create_session_with_config(config.clone()).await?;
528///     let session_id = session.id.clone();
529///     let created_at = session.created_at;
530///     
531///     // Later, restore the session
532///     let restored = manager.restore_session(session_id, config, created_at).await?;
533///     restored.start().await?;
534///     
535///     Ok(())
536/// }
537/// ```
538pub struct SessionManager {
539    /// Active sessions
540    sessions: Arc<DashMap<SessionId, Arc<AISession>>>,
541    /// Default session configuration
542    default_config: SessionConfig,
543}
544
545impl SessionManager {
546    /// Create a new session manager
547    pub fn new() -> Self {
548        Self {
549            sessions: Arc::new(DashMap::new()),
550            default_config: SessionConfig::default(),
551        }
552    }
553
554    /// Create a new session with default config
555    pub async fn create_session(&self) -> Result<Arc<AISession>> {
556        self.create_session_with_config(self.default_config.clone())
557            .await
558    }
559
560    /// Create a new session with custom config
561    pub async fn create_session_with_config(
562        &self,
563        config: SessionConfig,
564    ) -> Result<Arc<AISession>> {
565        let session = Arc::new(AISession::new(config).await?);
566        self.sessions.insert(session.id.clone(), session.clone());
567        Ok(session)
568    }
569
570    /// Restore a session with a specific ID (for persistence)
571    pub async fn restore_session(
572        &self,
573        id: SessionId,
574        config: SessionConfig,
575        created_at: DateTime<Utc>,
576    ) -> Result<Arc<AISession>> {
577        // Check if session already exists
578        if self.sessions.contains_key(&id) {
579            return Err(SessionError::AlreadyExists(id).into());
580        }
581
582        let session = Arc::new(AISession::new_with_id(id.clone(), config, created_at).await?);
583        self.sessions.insert(id, session.clone());
584        Ok(session)
585    }
586
587    /// Get a session by ID
588    pub fn get_session(&self, id: &SessionId) -> Option<Arc<AISession>> {
589        self.sessions.get(id).map(|entry| entry.clone())
590    }
591
592    /// List all active sessions
593    pub fn list_sessions(&self) -> Vec<SessionId> {
594        self.sessions
595            .iter()
596            .map(|entry| entry.key().clone())
597            .collect()
598    }
599
600    /// List all active session references
601    pub fn list_session_refs(&self) -> Vec<Arc<AISession>> {
602        self.sessions
603            .iter()
604            .map(|entry| entry.value().clone())
605            .collect()
606    }
607
608    /// Remove a session
609    pub async fn remove_session(&self, id: &SessionId) -> Result<()> {
610        if let Some((_, session)) = self.sessions.remove(id) {
611            session.stop().await?;
612        }
613        Ok(())
614    }
615
616    /// Clean up terminated sessions
617    pub async fn cleanup_terminated(&self) -> Result<usize> {
618        let mut removed = 0;
619        let terminated_ids: Vec<SessionId> = self
620            .sessions
621            .iter()
622            .filter(|entry| {
623                let session = entry.value();
624                if let Ok(status) = session.status.try_read() {
625                    *status == SessionStatus::Terminated
626                } else {
627                    false
628                }
629            })
630            .map(|entry| entry.key().clone())
631            .collect();
632
633        for id in terminated_ids {
634            self.sessions.remove(&id);
635            removed += 1;
636        }
637
638        Ok(removed)
639    }
640}
641
642impl Default for SessionManager {
643    fn default() -> Self {
644        Self::new()
645    }
646}
647
648/// Session error types
649
650#[cfg(test)]
651mod tests {
652    use super::*;
653
654    #[tokio::test]
655    async fn test_session_id() {
656        let id1 = SessionId::new();
657        let id2 = SessionId::new();
658        assert_ne!(id1, id2);
659    }
660
661    #[tokio::test]
662    async fn test_session_manager() {
663        let manager = SessionManager::new();
664        let session = manager.create_session().await.unwrap();
665
666        assert!(manager.get_session(&session.id).is_some());
667        assert_eq!(manager.list_sessions().len(), 1);
668
669        manager.remove_session(&session.id).await.unwrap();
670        assert!(manager.get_session(&session.id).is_none());
671    }
672}
ai_session/core/mod.rs

ai_session/core/
mod.rs
