//! Session persistence and storage management backed by SQLite use crate::types::Conversation; use crate::{Error, Result}; use aes_gcm::aead::{Aead, KeyInit}; use aes_gcm::{Aes256Gcm, Nonce}; use ring::rand::{SecureRandom, SystemRandom}; use serde::{Deserialize, Serialize}; use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous}; use sqlx::{Pool, Row, Sqlite}; use std::fs; use std::io::IsTerminal; use std::io::{self, Write}; use std::path::{Path, PathBuf}; use std::str::FromStr; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use uuid::Uuid; /// Metadata about a saved session #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SessionMeta { /// Conversation ID pub id: Uuid, /// Optional session name pub name: Option, /// Optional AI-generated description pub description: Option, /// Number of messages in the conversation pub message_count: usize, /// Model used pub model: String, /// When the session was created pub created_at: SystemTime, /// When the session was last updated pub updated_at: SystemTime, } /// Storage manager for persisting conversations in SQLite pub struct StorageManager { pool: Pool, database_path: PathBuf, } impl StorageManager { /// Create a new storage manager using the default database path pub async fn new() -> Result { let db_path = Self::default_database_path()?; Self::with_database_path(db_path).await } /// Create a storage manager using the provided database path pub async fn with_database_path(database_path: PathBuf) -> Result { if let Some(parent) = database_path.parent() { if !parent.exists() { std::fs::create_dir_all(parent).map_err(|e| { Error::Storage(format!( "Failed to create database directory {parent:?}: {e}" )) })?; } } let options = SqliteConnectOptions::from_str(&format!( "sqlite://{}", database_path .to_str() .ok_or_else(|| Error::Storage("Invalid database path".to_string()))? )) .map_err(|e| Error::Storage(format!("Invalid database URL: {e}")))? .create_if_missing(true) .journal_mode(SqliteJournalMode::Wal) .synchronous(SqliteSynchronous::Normal); let pool = SqlitePoolOptions::new() .max_connections(5) .connect_with(options) .await .map_err(|e| Error::Storage(format!("Failed to connect to database: {e}")))?; sqlx::migrate!("./migrations") .run(&pool) .await .map_err(|e| Error::Storage(format!("Failed to run database migrations: {e}")))?; let storage = Self { pool, database_path, }; storage.try_migrate_legacy_sessions().await?; Ok(storage) } /// Save a conversation. Existing entries are updated in-place. pub async fn save_conversation( &self, conversation: &Conversation, name: Option, ) -> Result<()> { self.save_conversation_with_description(conversation, name, None) .await } /// Save a conversation with an optional description override pub async fn save_conversation_with_description( &self, conversation: &Conversation, name: Option, description: Option, ) -> Result<()> { let mut serialized = conversation.clone(); if name.is_some() { serialized.name = name.clone(); } if description.is_some() { serialized.description = description.clone(); } let data = serde_json::to_string(&serialized) .map_err(|e| Error::Storage(format!("Failed to serialize conversation: {e}")))?; let created_at = to_epoch_seconds(serialized.created_at); let updated_at = to_epoch_seconds(serialized.updated_at); let message_count = serialized.messages.len() as i64; sqlx::query( r#" INSERT INTO conversations ( id, name, description, model, message_count, created_at, updated_at, data ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) ON CONFLICT(id) DO UPDATE SET name = excluded.name, description = excluded.description, model = excluded.model, message_count = excluded.message_count, created_at = excluded.created_at, updated_at = excluded.updated_at, data = excluded.data "#, ) .bind(serialized.id.to_string()) .bind(name.or(serialized.name.clone())) .bind(description.or(serialized.description.clone())) .bind(&serialized.model) .bind(message_count) .bind(created_at) .bind(updated_at) .bind(data) .execute(&self.pool) .await .map_err(|e| Error::Storage(format!("Failed to save conversation: {e}")))?; Ok(()) } /// Load a conversation by ID pub async fn load_conversation(&self, id: Uuid) -> Result { let record = sqlx::query(r#"SELECT data FROM conversations WHERE id = ?1"#) .bind(id.to_string()) .fetch_optional(&self.pool) .await .map_err(|e| Error::Storage(format!("Failed to load conversation: {e}")))?; let row = record.ok_or_else(|| Error::Storage(format!("No conversation found with id {id}")))?; let data: String = row .try_get("data") .map_err(|e| Error::Storage(format!("Failed to read conversation payload: {e}")))?; serde_json::from_str(&data) .map_err(|e| Error::Storage(format!("Failed to deserialize conversation: {e}"))) } /// List metadata for all saved conversations ordered by most recent update pub async fn list_sessions(&self) -> Result> { let rows = sqlx::query( r#" SELECT id, name, description, model, message_count, created_at, updated_at FROM conversations ORDER BY updated_at DESC "#, ) .fetch_all(&self.pool) .await .map_err(|e| Error::Storage(format!("Failed to list sessions: {e}")))?; let mut sessions = Vec::with_capacity(rows.len()); for row in rows { let id_text: String = row .try_get("id") .map_err(|e| Error::Storage(format!("Failed to read id column: {e}")))?; let id = Uuid::parse_str(&id_text) .map_err(|e| Error::Storage(format!("Invalid UUID in storage: {e}")))?; let message_count: i64 = row .try_get("message_count") .map_err(|e| Error::Storage(format!("Failed to read message count: {e}")))?; let created_at: i64 = row .try_get("created_at") .map_err(|e| Error::Storage(format!("Failed to read created_at: {e}")))?; let updated_at: i64 = row .try_get("updated_at") .map_err(|e| Error::Storage(format!("Failed to read updated_at: {e}")))?; sessions.push(SessionMeta { id, name: row .try_get("name") .map_err(|e| Error::Storage(format!("Failed to read name: {e}")))?, description: row .try_get("description") .map_err(|e| Error::Storage(format!("Failed to read description: {e}")))?, model: row .try_get("model") .map_err(|e| Error::Storage(format!("Failed to read model: {e}")))?, message_count: message_count as usize, created_at: from_epoch_seconds(created_at), updated_at: from_epoch_seconds(updated_at), }); } Ok(sessions) } /// Delete a conversation by ID pub async fn delete_session(&self, id: Uuid) -> Result<()> { sqlx::query("DELETE FROM conversations WHERE id = ?1") .bind(id.to_string()) .execute(&self.pool) .await .map_err(|e| Error::Storage(format!("Failed to delete conversation: {e}")))?; Ok(()) } pub async fn store_secure_item( &self, key: &str, plaintext: &[u8], master_key: &[u8], ) -> Result<()> { let cipher = create_cipher(master_key)?; let nonce_bytes = generate_nonce()?; let nonce = Nonce::from_slice(&nonce_bytes); let ciphertext = cipher .encrypt(nonce, plaintext) .map_err(|e| Error::Storage(format!("Failed to encrypt secure item: {e}")))?; let now = to_epoch_seconds(SystemTime::now()); sqlx::query( r#" INSERT INTO secure_items (key, nonce, ciphertext, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5) ON CONFLICT(key) DO UPDATE SET nonce = excluded.nonce, ciphertext = excluded.ciphertext, updated_at = excluded.updated_at "#, ) .bind(key) .bind(&nonce_bytes[..]) .bind(&ciphertext[..]) .bind(now) .bind(now) .execute(&self.pool) .await .map_err(|e| Error::Storage(format!("Failed to store secure item: {e}")))?; Ok(()) } pub async fn load_secure_item(&self, key: &str, master_key: &[u8]) -> Result>> { let record = sqlx::query("SELECT nonce, ciphertext FROM secure_items WHERE key = ?1") .bind(key) .fetch_optional(&self.pool) .await .map_err(|e| Error::Storage(format!("Failed to load secure item: {e}")))?; let Some(row) = record else { return Ok(None); }; let nonce_bytes: Vec = row .try_get("nonce") .map_err(|e| Error::Storage(format!("Failed to read secure item nonce: {e}")))?; let ciphertext: Vec = row .try_get("ciphertext") .map_err(|e| Error::Storage(format!("Failed to read secure item ciphertext: {e}")))?; if nonce_bytes.len() != 12 { return Err(Error::Storage( "Invalid nonce length for secure item".to_string(), )); } let cipher = create_cipher(master_key)?; let nonce = Nonce::from_slice(&nonce_bytes); let plaintext = cipher .decrypt(nonce, ciphertext.as_ref()) .map_err(|e| Error::Storage(format!("Failed to decrypt secure item: {e}")))?; Ok(Some(plaintext)) } pub async fn delete_secure_item(&self, key: &str) -> Result<()> { sqlx::query("DELETE FROM secure_items WHERE key = ?1") .bind(key) .execute(&self.pool) .await .map_err(|e| Error::Storage(format!("Failed to delete secure item: {e}")))?; Ok(()) } pub async fn clear_secure_items(&self) -> Result<()> { sqlx::query("DELETE FROM secure_items") .execute(&self.pool) .await .map_err(|e| Error::Storage(format!("Failed to clear secure items: {e}")))?; Ok(()) } /// Database location used by this storage manager pub fn database_path(&self) -> &Path { &self.database_path } /// Determine default database path (platform specific) pub fn default_database_path() -> Result { let data_dir = dirs::data_local_dir() .ok_or_else(|| Error::Storage("Could not determine data directory".to_string()))?; Ok(data_dir.join("owlen").join("owlen.db")) } fn legacy_sessions_dir() -> Result { let data_dir = dirs::data_local_dir() .ok_or_else(|| Error::Storage("Could not determine data directory".to_string()))?; Ok(data_dir.join("owlen").join("sessions")) } async fn database_has_records(&self) -> Result { let (count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM conversations") .fetch_one(&self.pool) .await .map_err(|e| Error::Storage(format!("Failed to inspect database: {e}")))?; Ok(count > 0) } async fn try_migrate_legacy_sessions(&self) -> Result<()> { if self.database_has_records().await? { return Ok(()); } let legacy_dir = match Self::legacy_sessions_dir() { Ok(dir) => dir, Err(_) => return Ok(()), }; if !legacy_dir.exists() { return Ok(()); } let entries = fs::read_dir(&legacy_dir).map_err(|e| { Error::Storage(format!("Failed to read legacy sessions directory: {e}")) })?; let mut json_files = Vec::new(); for entry in entries.flatten() { let path = entry.path(); if path.extension().and_then(|s| s.to_str()) == Some("json") { json_files.push(path); } } if json_files.is_empty() { return Ok(()); } if !io::stdin().is_terminal() { return Ok(()); } println!( "Legacy OWLEN session files were found in {}.", legacy_dir.display() ); if !prompt_yes_no("Migrate them to the new SQLite storage? (y/N) ")? { println!("Skipping legacy session migration."); return Ok(()); } println!("Migrating legacy sessions..."); let mut migrated = 0usize; for path in &json_files { match fs::read_to_string(path) { Ok(content) => match serde_json::from_str::(&content) { Ok(conversation) => { if let Err(err) = self .save_conversation_with_description( &conversation, conversation.name.clone(), conversation.description.clone(), ) .await { println!(" • Failed to migrate {}: {}", path.display(), err); } else { migrated += 1; } } Err(err) => { println!( " • Failed to parse conversation {}: {}", path.display(), err ); } }, Err(err) => { println!(" • Failed to read {}: {}", path.display(), err); } } } if migrated > 0 { if let Err(err) = archive_legacy_directory(&legacy_dir) { println!( "Warning: migrated sessions but failed to archive legacy directory: {}", err ); } } println!("Migrated {} legacy sessions.", migrated); Ok(()) } } fn to_epoch_seconds(time: SystemTime) -> i64 { match time.duration_since(UNIX_EPOCH) { Ok(duration) => duration.as_secs() as i64, Err(_) => 0, } } fn from_epoch_seconds(seconds: i64) -> SystemTime { UNIX_EPOCH + Duration::from_secs(seconds.max(0) as u64) } fn prompt_yes_no(prompt: &str) -> Result { print!("{}", prompt); io::stdout() .flush() .map_err(|e| Error::Storage(format!("Failed to flush stdout: {e}")))?; let mut input = String::new(); io::stdin() .read_line(&mut input) .map_err(|e| Error::Storage(format!("Failed to read input: {e}")))?; let trimmed = input.trim().to_lowercase(); Ok(matches!(trimmed.as_str(), "y" | "yes")) } fn archive_legacy_directory(legacy_dir: &Path) -> Result<()> { let mut backup_dir = legacy_dir.with_file_name("sessions_legacy_backup"); let mut counter = 1; while backup_dir.exists() { backup_dir = legacy_dir.with_file_name(format!("sessions_legacy_backup_{}", counter)); counter += 1; } fs::rename(legacy_dir, &backup_dir).map_err(|e| { Error::Storage(format!( "Failed to archive legacy sessions directory {}: {}", legacy_dir.display(), e )) })?; println!("Legacy session files archived to {}", backup_dir.display()); Ok(()) } fn create_cipher(master_key: &[u8]) -> Result { if master_key.len() != 32 { return Err(Error::Storage( "Master key must be 32 bytes for AES-256-GCM".to_string(), )); } Aes256Gcm::new_from_slice(master_key).map_err(|_| { Error::Storage("Failed to initialize cipher with provided master key".to_string()) }) } fn generate_nonce() -> Result<[u8; 12]> { let mut nonce = [0u8; 12]; SystemRandom::new() .fill(&mut nonce) .map_err(|_| Error::Storage("Failed to generate nonce".to_string()))?; Ok(nonce) } #[cfg(test)] mod tests { use super::*; use crate::types::{Conversation, Message}; use tempfile::tempdir; fn sample_conversation() -> Conversation { Conversation { id: Uuid::new_v4(), name: Some("Test conversation".to_string()), description: Some("A sample conversation".to_string()), messages: vec![ Message::user("Hello".to_string()), Message::assistant("Hi".to_string()), ], model: "test-model".to_string(), created_at: SystemTime::now(), updated_at: SystemTime::now(), } } #[tokio::test] async fn test_storage_lifecycle() { let temp_dir = tempdir().expect("failed to create temp dir"); let db_path = temp_dir.path().join("owlen.db"); let storage = StorageManager::with_database_path(db_path).await.unwrap(); let conversation = sample_conversation(); storage .save_conversation(&conversation, None) .await .expect("failed to save conversation"); let sessions = storage.list_sessions().await.unwrap(); assert_eq!(sessions.len(), 1); assert_eq!(sessions[0].id, conversation.id); let loaded = storage.load_conversation(conversation.id).await.unwrap(); assert_eq!(loaded.messages.len(), 2); storage .delete_session(conversation.id) .await .expect("failed to delete conversation"); let sessions = storage.list_sessions().await.unwrap(); assert!(sessions.is_empty()); } }