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 attention;
85pub mod headless;
86pub mod lifecycle;
87pub mod process;
88pub mod pty;
89pub mod terminal;
90
91pub use attention::AttentionState;
92use attention::AttentionTracker;
93
94use crate::context::SessionContext;
95use crate::persistence::CommandRecord;
96
97/// Session error type
98#[derive(Debug, thiserror::Error)]
99pub enum SessionError {
100    #[error("Session not found: {0}")]
101    NotFound(SessionId),
102
103    #[error("Session already exists: {0}")]
104    AlreadyExists(SessionId),
105
106    #[error("PTY error: {0}")]
107    PtyError(String),
108
109    #[error("Process error: {0}")]
110    ProcessError(String),
111
112    #[error("IO error: {0}")]
113    IoError(#[from] std::io::Error),
114
115    #[error("Other error: {0}")]
116    Other(#[from] anyhow::Error),
117}
118
119/// Session result type
120pub type SessionResult<T> = std::result::Result<T, SessionError>;
121
122/// Unique session identifier
123#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)]
124pub struct SessionId(Uuid);
125
126impl Default for SessionId {
127    fn default() -> Self {
128        Self::new()
129    }
130}
131
132impl SessionId {
133    /// Create a new unique session ID
134    pub fn new() -> Self {
135        Self(Uuid::new_v4())
136    }
137
138    /// Create a new v4 UUID session ID
139    pub fn new_v4() -> Self {
140        Self(Uuid::new_v4())
141    }
142
143    /// Parse from string
144    pub fn parse_str(s: &str) -> Result<Self> {
145        Ok(Self(Uuid::parse_str(s)?))
146    }
147
148    /// Convert to string
149    pub fn to_string(&self) -> String {
150        self.0.to_string()
151    }
152
153    /// Get inner UUID
154    pub fn as_uuid(&self) -> &Uuid {
155        &self.0
156    }
157}
158
159impl std::fmt::Display for SessionId {
160    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
161        write!(f, "{}", self.0)
162    }
163}
164
165/// Session status
166#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
167pub enum SessionStatus {
168    /// Session is being initialized
169    #[default]
170    Initializing,
171    /// Session is running and ready
172    Running,
173    /// Session is paused
174    Paused,
175    /// Session is being terminated
176    Terminating,
177    /// Session has been terminated
178    Terminated,
179    /// Session encountered an error
180    Error,
181}
182
183/// Session configuration
184#[derive(Debug, Clone, Serialize, Deserialize)]
185#[serde(default)]
186pub struct SessionConfig {
187    /// Session name (optional)
188    pub name: Option<String>,
189    /// Working directory
190    pub working_directory: PathBuf,
191    /// Environment variables
192    pub environment: HashMap<String, String>,
193    /// Shell command to execute
194    pub shell: Option<String>,
195    /// Shell command to execute (alternative field for compatibility)
196    pub shell_command: Option<String>,
197    /// PTY size (rows, cols)
198    pub pty_size: (u16, u16),
199    /// Output buffer size in bytes
200    pub output_buffer_size: usize,
201    /// Session timeout (None for no timeout)
202    pub timeout: Option<Duration>,
203    /// Enable output compression
204    pub compress_output: bool,
205    /// Enable semantic output parsing
206    pub parse_output: bool,
207    /// Enable AI features
208    pub enable_ai_features: bool,
209    /// Context configuration
210    pub context_config: ContextConfig,
211    /// Agent role (optional)
212    pub agent_role: Option<String>,
213    /// Force headless (non-PTY) execution (useful for restricted sandboxes)
214    pub force_headless: bool,
215    /// Allow automatic fallback to headless mode when PTY creation fails
216    pub allow_headless_fallback: bool,
217}
218
219/// Context configuration for AI features
220#[derive(Debug, Clone, Serialize, Deserialize)]
221#[serde(default)]
222pub struct ContextConfig {
223    /// Maximum tokens for context
224    pub max_tokens: usize,
225    /// Compression threshold (0.0 to 1.0)
226    pub compression_threshold: f64,
227}
228
229impl Default for SessionConfig {
230    fn default() -> Self {
231        Self {
232            name: None,
233            working_directory: std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/")),
234            environment: HashMap::new(),
235            shell: None,
236            shell_command: None,
237            pty_size: (24, 80),
238            output_buffer_size: 1024 * 1024, // 1MB
239            timeout: None,
240            compress_output: true,
241            parse_output: true,
242            enable_ai_features: false,
243            context_config: ContextConfig::default(),
244            agent_role: None,
245            force_headless: false,
246            allow_headless_fallback: true,
247        }
248    }
249}
250
251impl Default for ContextConfig {
252    fn default() -> Self {
253        Self {
254            max_tokens: 4096,
255            compression_threshold: 0.8,
256        }
257    }
258}
259
260/// AI-optimized session
261pub struct AISession {
262    /// Unique session ID
263    pub id: SessionId,
264    /// Session configuration
265    pub config: SessionConfig,
266    /// Current status
267    pub status: RwLock<SessionStatus>,
268    /// Session context (AI state, history, etc.)
269    pub context: Arc<RwLock<SessionContext>>,
270    /// Process handle
271    process: Arc<RwLock<Option<process::ProcessHandle>>>,
272    /// Terminal handle (PTY or headless)
273    terminal: Arc<RwLock<Option<terminal::TerminalHandle>>>,
274    /// Creation time
275    pub created_at: DateTime<Utc>,
276    /// Last activity time
277    pub last_activity: Arc<RwLock<DateTime<Utc>>>,
278    /// Session metadata
279    pub metadata: Arc<RwLock<HashMap<String, serde_json::Value>>>,
280    /// Command history tracking
281    pub command_history: Arc<RwLock<Vec<CommandRecord>>>,
282    /// Command count
283    pub command_count: Arc<RwLock<usize>>,
284    /// Total tokens used
285    pub total_tokens: Arc<RwLock<usize>>,
286    /// Attention state — what does this session need next?
287    attention: AttentionTracker,
288}
289
290impl AISession {
291    /// Create a new AI session
292    pub async fn new(config: SessionConfig) -> Result<Self> {
293        let id = SessionId::new();
294        let now = Utc::now();
295
296        Ok(Self {
297            id: id.clone(),
298            config,
299            status: RwLock::new(SessionStatus::Initializing),
300            context: Arc::new(RwLock::new(SessionContext::new(id))),
301            process: Arc::new(RwLock::new(None)),
302            terminal: Arc::new(RwLock::new(None)),
303            created_at: now,
304            last_activity: Arc::new(RwLock::new(now)),
305            metadata: Arc::new(RwLock::new(HashMap::new())),
306            command_history: Arc::new(RwLock::new(Vec::new())),
307            command_count: Arc::new(RwLock::new(0)),
308            total_tokens: Arc::new(RwLock::new(0)),
309            attention: AttentionTracker::new(AttentionState::Idle),
310        })
311    }
312
313    /// Create an AI session with a specific ID (for restoration)
314    pub async fn new_with_id(
315        id: SessionId,
316        config: SessionConfig,
317        created_at: DateTime<Utc>,
318    ) -> Result<Self> {
319        let now = Utc::now();
320
321        Ok(Self {
322            id: id.clone(),
323            config,
324            status: RwLock::new(SessionStatus::Initializing),
325            context: Arc::new(RwLock::new(SessionContext::new(id))),
326            process: Arc::new(RwLock::new(None)),
327            terminal: Arc::new(RwLock::new(None)),
328            created_at,
329            last_activity: Arc::new(RwLock::new(now)),
330            metadata: Arc::new(RwLock::new(HashMap::new())),
331            command_history: Arc::new(RwLock::new(Vec::new())),
332            command_count: Arc::new(RwLock::new(0)),
333            total_tokens: Arc::new(RwLock::new(0)),
334            attention: AttentionTracker::new(AttentionState::Idle),
335        })
336    }
337
338    /// Start the session
339    pub async fn start(&self) -> Result<()> {
340        lifecycle::start_session(self).await
341    }
342
343    /// Stop the session
344    pub async fn stop(&self) -> Result<()> {
345        lifecycle::stop_session(self).await
346    }
347
348    /// Send input to the session
349    pub async fn send_input(&self, input: &str) -> Result<()> {
350        let terminal_guard = self.terminal.read().await;
351        if let Some(terminal) = terminal_guard.as_ref() {
352            terminal.write(input.as_bytes()).await?;
353            *self.last_activity.write().await = Utc::now();
354            Ok(())
355        } else {
356            Err(anyhow::anyhow!("Session not started"))
357        }
358    }
359
360    /// Read output from the session
361    pub async fn read_output(&self) -> Result<Vec<u8>> {
362        let terminal = self.terminal.read().await;
363        if let Some(terminal) = terminal.as_ref() {
364            let output = terminal.read().await?;
365            *self.last_activity.write().await = Utc::now();
366            Ok(output)
367        } else {
368            Err(anyhow::anyhow!("Session not started"))
369        }
370    }
371
372    /// Get current session status
373    pub async fn status(&self) -> SessionStatus {
374        *self.status.read().await
375    }
376
377    /// Update session metadata
378    pub async fn set_metadata(&self, key: String, value: serde_json::Value) -> Result<()> {
379        self.metadata.write().await.insert(key, value);
380        Ok(())
381    }
382
383    /// Get session metadata
384    pub async fn get_metadata(&self, key: &str) -> Option<serde_json::Value> {
385        self.metadata.read().await.get(key).cloned()
386    }
387
388    /// Execute a command and record it in history
389    pub async fn execute_command(&self, command: &str) -> Result<String> {
390        let start_time = Utc::now();
391
392        let shell_env = std::env::var("SHELL").ok();
393        let shell = self
394            .config
395            .shell
396            .as_deref()
397            .or(shell_env.as_deref())
398            .unwrap_or("/bin/sh");
399        let mut cmd = tokio::process::Command::new(shell);
400        cmd.arg("-lc")
401            .arg(command)
402            .current_dir(&self.config.working_directory);
403        for (key, value) in &self.config.environment {
404            cmd.env(key, value);
405        }
406
407        let execution =
408            crate::execution::run_provider_command(cmd, &self.config.working_directory, shell)
409                .await?;
410        *self.last_activity.write().await = Utc::now();
411        let output = if execution.stderr.is_empty() {
412            execution.stdout
413        } else if execution.stdout.is_empty() {
414            execution.stderr
415        } else {
416            format!("{}{}", execution.stdout, execution.stderr)
417        };
418
419        // Record the command in history
420        let end_time = Utc::now();
421        let duration_ms = execution
422            .duration_ms
423            .max((end_time - start_time).num_milliseconds() as u64);
424
425        let record = CommandRecord {
426            command: command.to_string(),
427            timestamp: start_time,
428            exit_code: execution.status.code(),
429            output_preview: if output.len() > 200 {
430                format!("{}...", &output[..200])
431            } else {
432                output.clone()
433            },
434            duration_ms,
435        };
436
437        // Update history and counters
438        self.command_history.write().await.push(record);
439        *self.command_count.write().await += 1;
440
441        Ok(output)
442    }
443
444    /// Add tokens to the session total
445    pub async fn add_tokens(&self, token_count: usize) {
446        *self.total_tokens.write().await += token_count;
447    }
448
449    /// Get command history
450    pub async fn get_command_history(&self) -> Vec<CommandRecord> {
451        self.command_history.read().await.clone()
452    }
453
454    /// Get command count
455    pub async fn get_command_count(&self) -> usize {
456        *self.command_count.read().await
457    }
458
459    /// Get total tokens used
460    pub async fn get_total_tokens(&self) -> usize {
461        *self.total_tokens.read().await
462    }
463
464    /// Current attention state — what triage bucket this session falls into.
465    pub fn attention(&self) -> AttentionState {
466        self.attention.get()
467    }
468
469    /// Subscribe to attention state changes.
470    ///
471    /// Returns a `watch::Receiver` initialized to the current state. Each
472    /// subsequent transition wakes `.changed()`.
473    pub fn subscribe_attention(&self) -> tokio::sync::watch::Receiver<AttentionState> {
474        self.attention.subscribe()
475    }
476
477    /// Set the attention state explicitly. Returns `true` if the state changed.
478    pub fn set_attention(&self, state: AttentionState) -> bool {
479        self.attention.set(state)
480    }
481
482    /// Update attention from a parsed output snapshot, if the snapshot is
483    /// decisive enough to imply a transition. Returns the new state when one
484    /// was applied.
485    pub fn update_attention_from_parsed(
486        &self,
487        parsed: &crate::output::ParsedOutput,
488    ) -> Option<AttentionState> {
489        let next = AttentionState::from_parsed(parsed)?;
490        self.attention.set(next);
491        Some(next)
492    }
493
494    /// Clear command history (keep recent N commands)
495    pub async fn trim_command_history(&self, keep_recent: usize) {
496        let mut history = self.command_history.write().await;
497        if history.len() > keep_recent {
498            let start_index = history.len() - keep_recent;
499            history.drain(0..start_index);
500        }
501    }
502}
503
504/// AI-optimized session manager for creating and managing multiple terminal sessions.
505///
506/// The `SessionManager` provides a centralized way to create, track, and manage AI-enhanced
507/// terminal sessions. It includes automatic cleanup, session restoration, and efficient
508/// resource management.
509///
510/// # Features
511///
512/// - **Session Pooling**: Efficient management of multiple concurrent sessions
513/// - **Automatic Cleanup**: Garbage collection of terminated sessions
514/// - **Session Restoration**: Restore sessions from persistent storage
515/// - **Resource Management**: Automatic cleanup and memory management
516///
517/// # Examples
518///
519/// ## Basic Session Management
520///
521/// ```no_run
522/// use ai_session::{SessionManager, SessionConfig};
523///
524/// #[tokio::main]
525/// async fn main() -> anyhow::Result<()> {
526///     let manager = SessionManager::new();
527///     
528///     // Create multiple sessions
529///     let session1 = manager.create_session().await?;
530///     let session2 = manager.create_session().await?;
531///     
532///     session1.start().await?;
533///     session2.start().await?;
534///     
535///     // List all active sessions
536///     let session_ids = manager.list_sessions();
537///     println!("Active sessions: {}", session_ids.len());
538///     
539///     // Clean up
540///     manager.remove_session(&session1.id).await?;
541///     manager.remove_session(&session2.id).await?;
542///     
543///     Ok(())
544/// }
545/// ```
546///
547/// ## Custom Configuration
548///
549/// ```no_run
550/// use ai_session::{SessionManager, SessionConfig, ContextConfig};
551///
552/// #[tokio::main]
553/// async fn main() -> anyhow::Result<()> {
554///     let manager = SessionManager::new();
555///     
556///     // Configure for AI development agent
557///     let mut config = SessionConfig::default();
558///     config.enable_ai_features = true;
559///     config.agent_role = Some("backend-developer".to_string());
560///     config.working_directory = "/project/backend".into();
561///     config.context_config = ContextConfig {
562///         max_tokens: 8192,
563///         compression_threshold: 0.8,
564///     };
565///     
566///     let session = manager.create_session_with_config(config).await?;
567///     session.start().await?;
568///     
569///     // Session is optimized for AI backend development
570///     Ok(())
571/// }
572/// ```
573///
574/// ## Session Persistence
575///
576/// ```no_run
577/// use ai_session::{SessionManager, SessionConfig, SessionId};
578/// use chrono::Utc;
579///
580/// #[tokio::main]
581/// async fn main() -> anyhow::Result<()> {
582///     let manager = SessionManager::new();
583///     
584///     // Create session
585///     let config = SessionConfig::default();
586///     let session = manager.create_session_with_config(config.clone()).await?;
587///     let session_id = session.id.clone();
588///     let created_at = session.created_at;
589///     
590///     // Later, restore the session
591///     let restored = manager.restore_session(session_id, config, created_at).await?;
592///     restored.start().await?;
593///     
594///     Ok(())
595/// }
596/// ```
597pub struct SessionManager {
598    /// Active sessions
599    sessions: Arc<DashMap<SessionId, Arc<AISession>>>,
600    /// Default session configuration
601    default_config: SessionConfig,
602}
603
604impl SessionManager {
605    /// Create a new session manager
606    pub fn new() -> Self {
607        Self {
608            sessions: Arc::new(DashMap::new()),
609            default_config: SessionConfig::default(),
610        }
611    }
612
613    /// Create a new session with default config
614    pub async fn create_session(&self) -> Result<Arc<AISession>> {
615        self.create_session_with_config(self.default_config.clone())
616            .await
617    }
618
619    /// Create a new session with custom config
620    pub async fn create_session_with_config(
621        &self,
622        config: SessionConfig,
623    ) -> Result<Arc<AISession>> {
624        let session = Arc::new(AISession::new(config).await?);
625        self.sessions.insert(session.id.clone(), session.clone());
626        Ok(session)
627    }
628
629    /// Restore a session with a specific ID (for persistence)
630    pub async fn restore_session(
631        &self,
632        id: SessionId,
633        config: SessionConfig,
634        created_at: DateTime<Utc>,
635    ) -> Result<Arc<AISession>> {
636        // Check if session already exists
637        if self.sessions.contains_key(&id) {
638            return Err(SessionError::AlreadyExists(id).into());
639        }
640
641        let session = Arc::new(AISession::new_with_id(id.clone(), config, created_at).await?);
642        self.sessions.insert(id, session.clone());
643        Ok(session)
644    }
645
646    /// Get a session by ID
647    pub fn get_session(&self, id: &SessionId) -> Option<Arc<AISession>> {
648        self.sessions.get(id).map(|entry| entry.clone())
649    }
650
651    /// List all active sessions
652    pub fn list_sessions(&self) -> Vec<SessionId> {
653        self.sessions
654            .iter()
655            .map(|entry| entry.key().clone())
656            .collect()
657    }
658
659    /// List all active session references
660    pub fn list_session_refs(&self) -> Vec<Arc<AISession>> {
661        self.sessions
662            .iter()
663            .map(|entry| entry.value().clone())
664            .collect()
665    }
666
667    /// Remove a session
668    pub async fn remove_session(&self, id: &SessionId) -> Result<()> {
669        if let Some((_, session)) = self.sessions.remove(id) {
670            session.stop().await?;
671        }
672        Ok(())
673    }
674
675    /// Clean up terminated sessions
676    pub async fn cleanup_terminated(&self) -> Result<usize> {
677        let mut removed = 0;
678        let terminated_ids: Vec<SessionId> = self
679            .sessions
680            .iter()
681            .filter(|entry| {
682                let session = entry.value();
683                if let Ok(status) = session.status.try_read() {
684                    *status == SessionStatus::Terminated
685                } else {
686                    false
687                }
688            })
689            .map(|entry| entry.key().clone())
690            .collect();
691
692        for id in terminated_ids {
693            self.sessions.remove(&id);
694            removed += 1;
695        }
696
697        Ok(removed)
698    }
699
700    /// Find a session by its name (exact match)
701    ///
702    /// # Arguments
703    /// * `name` - The exact name to search for
704    ///
705    /// # Returns
706    /// The session if found, None otherwise
707    pub fn find_session_by_name(&self, name: &str) -> Option<Arc<AISession>> {
708        self.sessions.iter().find_map(|entry| {
709            let session = entry.value();
710            if session.config.name.as_deref() == Some(name) {
711                Some(session.clone())
712            } else {
713                None
714            }
715        })
716    }
717
718    /// Find sessions with names matching a prefix
719    ///
720    /// # Arguments
721    /// * `prefix` - The prefix to match against session names
722    ///
723    /// # Returns
724    /// List of sessions whose names start with the given prefix
725    pub fn list_sessions_by_prefix(&self, prefix: &str) -> Vec<Arc<AISession>> {
726        self.sessions
727            .iter()
728            .filter_map(|entry| {
729                let session = entry.value();
730                if session
731                    .config
732                    .name
733                    .as_ref()
734                    .map(|n| n.starts_with(prefix))
735                    .unwrap_or(false)
736                {
737                    Some(session.clone())
738                } else {
739                    None
740                }
741            })
742            .collect()
743    }
744
745    /// List all sessions with detailed information (snapshot, no locking required)
746    ///
747    /// Returns a vector of SessionInfo containing summary information about all active sessions.
748    /// This is more efficient than listing sessions and then querying each one individually.
749    pub fn list_sessions_detailed(&self) -> Vec<crate::SessionInfo> {
750        self.sessions
751            .iter()
752            .map(|entry| {
753                let session = entry.value();
754                let status = session
755                    .status
756                    .try_read()
757                    .map(|s| *s)
758                    .unwrap_or(SessionStatus::Initializing);
759                let last_activity = session
760                    .last_activity
761                    .try_read()
762                    .map(|t| *t)
763                    .unwrap_or(session.created_at);
764                let command_count = session.command_count.try_read().map(|c| *c).unwrap_or(0);
765                let context_tokens = session.total_tokens.try_read().map(|t| *t).unwrap_or(0);
766
767                crate::SessionInfo {
768                    id: session.id.clone(),
769                    name: session.config.name.clone(),
770                    status,
771                    created_at: session.created_at,
772                    last_activity,
773                    working_directory: session.config.working_directory.clone(),
774                    ai_features_enabled: session.config.enable_ai_features,
775                    context_token_count: context_tokens,
776                    command_count,
777                }
778            })
779            .collect()
780    }
781}
782
783impl Default for SessionManager {
784    fn default() -> Self {
785        Self::new()
786    }
787}
788
789/// Session error types
790#[cfg(test)]
791mod tests {
792    use super::*;
793
794    #[tokio::test]
795    async fn test_session_id() {
796        let id1 = SessionId::new();
797        let id2 = SessionId::new();
798        assert_ne!(id1, id2);
799    }
800
801    #[tokio::test]
802    async fn test_session_manager() {
803        let manager = SessionManager::new();
804        let session = manager.create_session().await.unwrap();
805
806        assert!(manager.get_session(&session.id).is_some());
807        assert_eq!(manager.list_sessions().len(), 1);
808
809        manager.remove_session(&session.id).await.unwrap();
810        assert!(manager.get_session(&session.id).is_none());
811    }
812
813    #[tokio::test]
814    async fn execute_command_runs_in_session_working_directory() {
815        let dir = tempfile::tempdir().unwrap();
816        let config = SessionConfig {
817            working_directory: dir.path().to_path_buf(),
818            shell: Some("/bin/sh".to_string()),
819            ..SessionConfig::default()
820        };
821        let session = AISession::new(config).await.unwrap();
822
823        let output = session
824            .execute_command("printf ai-session-ok && pwd")
825            .await
826            .unwrap();
827
828        assert!(output.contains("ai-session-ok"));
829        assert!(output.contains(&dir.path().to_string_lossy().to_string()));
830
831        let history = session.get_command_history().await;
832        assert_eq!(history.len(), 1);
833        assert_eq!(history[0].exit_code, Some(0));
834    }
835}
ai_session/core/mod.rs

ai_session/core/
mod.rs
