ai_session/
session_persistence.rs

1//! Persistent session manager for the CLI
2//!
3//! This module provides a SessionManager wrapper that persists sessions
4//! across CLI invocations using the ai-session library's PersistenceManager.
5
6use crate::{SessionConfig, SessionManager as InnerSessionManager, AISession};
7use crate::core::{SessionId, SessionStatus};
8use crate::persistence::{PersistenceManager, SessionState, SessionMetadata};
9use anyhow::Result;
10use std::path::PathBuf;
11use std::sync::Arc;
12
13/// Global session manager instance
14static MANAGER_INSTANCE: tokio::sync::OnceCell<Arc<PersistentSessionManager>> = tokio::sync::OnceCell::const_new();
15
16/// Get the global persistent session manager
17pub async fn get_session_manager() -> Result<Arc<PersistentSessionManager>> {
18    Ok(MANAGER_INSTANCE
19        .get_or_init(|| async {
20            Arc::new(PersistentSessionManager::new().await.unwrap())
21        })
22        .await
23        .clone())
24}
25
26/// Persistent session manager that wraps the inner SessionManager
27/// and provides persistence across CLI invocations
28pub struct PersistentSessionManager {
29    /// Inner session manager
30    inner: InnerSessionManager,
31    /// Persistence manager for saving/loading sessions
32    persistence: PersistenceManager,
33    /// Cache directory for session data
34    #[allow(dead_code)]
35    storage_path: PathBuf,
36}
37
38impl PersistentSessionManager {
39    /// Create a new persistent session manager
40    pub async fn new() -> Result<Self> {
41        let storage_path = dirs::data_dir()
42            .ok_or_else(|| anyhow::anyhow!("Could not determine data directory"))?
43            .join("ai-session")
44            .join("sessions");
45        
46        std::fs::create_dir_all(&storage_path)?;
47        eprintln!("Using storage path: {}", storage_path.display());
48        
49        let persistence = PersistenceManager::new(storage_path.clone());
50        let inner = InnerSessionManager::new();
51        
52        let mut manager = Self {
53            inner,
54            persistence,
55            storage_path,
56        };
57        
58        // Load existing sessions
59        manager.restore_sessions().await?;
60        
61        Ok(manager)
62    }
63    
64    /// Restore sessions from persistent storage
65    async fn restore_sessions(&mut self) -> Result<()> {
66        eprintln!("Restoring sessions from persistence...");
67        let session_ids = self.persistence.list_sessions().await?;
68        eprintln!("Found {} sessions to restore", session_ids.len());
69        
70        for session_id in session_ids {
71            match self.persistence.load_session(&session_id).await {
72                Ok(state) => {
73                    // Only restore sessions that were running or paused
74                    match state.status {
75                        SessionStatus::Running | SessionStatus::Paused => {
76                            // Restore the session with its original ID
77                            match self.inner.restore_session(
78                                state.session_id.clone(),
79                                state.config.clone(),
80                                state.metadata.created_at,
81                            ).await {
82                                Ok(_session) => {
83                                    eprintln!("Restored session: {}", session_id);
84                                }
85                                Err(e) => {
86                                    eprintln!("Warning: Failed to restore session {}: {}", session_id, e);
87                                }
88                            }
89                        }
90                        _ => {
91                            // Skip terminated or error sessions
92                            eprintln!("Skipping terminated session: {}", session_id);
93                        }
94                    }
95                }
96                Err(e) => {
97                    eprintln!("Warning: Failed to restore session {}: {}", session_id, e);
98                }
99            }
100        }
101        
102        Ok(())
103    }
104    
105    /// Create a new session with config and persist it
106    pub async fn create_session_with_config(&self, config: SessionConfig) -> Result<Arc<AISession>> {
107        let session = self.inner.create_session_with_config(config).await?;
108        
109        // Start the session first
110        session.start().await?;
111        
112        // Save the session state after starting
113        let state = SessionState {
114            session_id: session.id.clone(),
115            config: session.config.clone(),
116            status: session.status().await,
117            context: session.context.read().await.clone(),
118            command_history: vec![],
119            metadata: SessionMetadata::default(),
120        };
121        
122        eprintln!("Saving session {} with status {:?}", session.id, state.status);
123        self.persistence.save_session(&session.id, &state).await?;
124        eprintln!("Session {} saved successfully", session.id);
125        
126        Ok(session)
127    }
128    
129    /// Get a session by ID, starting it if necessary
130    pub async fn get_session(&self, id: &SessionId) -> Option<Arc<AISession>> {
131        if let Some(session) = self.inner.get_session(id) {
132            // Check if session needs to be started
133            if session.status().await == SessionStatus::Initializing {
134                if let Err(e) = session.start().await {
135                    eprintln!("Error: Failed to start session {}: {}", id, e);
136                    return None;
137                }
138            }
139            Some(session)
140        } else {
141            // Try to load from persistence
142            match self.persistence.load_session(id).await {
143                Ok(state) => {
144                    // Recreate the session
145                    match self.inner.create_session_with_config(state.config.clone()).await {
146                        Ok(session) => {
147                            // Start the session
148                            if let Err(e) = session.start().await {
149                                eprintln!("Error: Failed to start restored session {}: {}", id, e);
150                                return None;
151                            }
152                            Some(session)
153                        }
154                        Err(e) => {
155                            eprintln!("Error: Failed to recreate session {}: {}", id, e);
156                            None
157                        }
158                    }
159                }
160                Err(e) => {
161                    eprintln!("Debug: Session {} not found in persistence: {}", id, e);
162                    None
163                }
164            }
165        }
166    }
167    
168    /// List all sessions (both active and persisted)
169    pub async fn list_all_sessions(&self) -> Result<Vec<SessionId>> {
170        let mut all_sessions = self.inner.list_sessions();
171        let persisted = self.persistence.list_sessions().await?;
172        
173        // Add persisted sessions that aren't already active
174        for session_id in persisted {
175            if !all_sessions.contains(&session_id) {
176                all_sessions.push(session_id);
177            }
178        }
179        
180        Ok(all_sessions)
181    }
182    
183    /// Remove a session and delete its persistent data
184    pub async fn remove_session(&self, id: &SessionId) -> Result<()> {
185        // Remove from active sessions
186        self.inner.remove_session(id).await?;
187        
188        // Delete persistent data
189        self.persistence.delete_session(id).await?;
190        
191        Ok(())
192    }
193    
194    /// Update session state in persistence
195    pub async fn update_session_state(&self, session: &AISession) -> Result<()> {
196        let state = SessionState {
197            session_id: session.id.clone(),
198            config: session.config.clone(),
199            status: session.status().await,
200            context: session.context.read().await.clone(),
201            command_history: vec![], // TODO: Implement command history tracking
202            metadata: SessionMetadata {
203                created_at: session.created_at,
204                last_accessed: *session.last_activity.read().await,
205                command_count: 0, // TODO: Track command count
206                total_tokens: 0, // TODO: Track token usage
207                custom: session.metadata.read().await.clone(),
208            },
209        };
210        
211        self.persistence.save_session(&session.id, &state).await?;
212        Ok(())
213    }
214}
ai_session/session_persistence.rs

ai_session/
session_persistence.rs
