use rusqlite::{params, Connection, Result as SqlResult}; use std::path::Path; use std::sync::Mutex; use tracing::{info, warn}; /// SQLite-backed persistent state for Sol. /// /// Stores: /// - Conversation registry: room_id → Mistral conversation_id + token estimates /// - Agent registry: agent_name → Mistral agent_id /// /// ## Kubernetes mount /// /// The database file should be on a persistent volume. Recommended setup: /// /// ```yaml /// # PersistentVolumeClaim (Longhorn) /// apiVersion: v1 /// kind: PersistentVolumeClaim /// metadata: /// name: sol-data /// namespace: matrix /// spec: /// accessModes: [ReadWriteOnce] /// storageClassName: longhorn /// resources: /// requests: /// storage: 1Gi /// /// # Deployment volume mount /// volumes: /// - name: sol-data /// persistentVolumeClaim: /// claimName: sol-data /// containers: /// - name: sol /// volumeMounts: /// - name: sol-data /// mountPath: /data /// ``` /// /// Default path: `/data/sol.db` (configurable via `matrix.db_path` in sol.toml). /// The `/data` mount also holds the Matrix SDK state store at `/data/matrix-state`. pub struct Store { conn: Mutex, } impl Store { /// Open or create the database at the given path. /// Creates tables if they don't exist. pub fn open(path: &str) -> anyhow::Result { // Ensure parent directory exists if let Some(parent) = Path::new(path).parent() { std::fs::create_dir_all(parent)?; } let conn = Connection::open(path)?; // Enable WAL mode for better concurrent read performance conn.execute_batch("PRAGMA journal_mode=WAL;")?; conn.execute_batch( "CREATE TABLE IF NOT EXISTS conversations ( room_id TEXT PRIMARY KEY, conversation_id TEXT NOT NULL, estimated_tokens INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL DEFAULT (datetime('now')) ); CREATE TABLE IF NOT EXISTS agents ( name TEXT PRIMARY KEY, agent_id TEXT NOT NULL, model TEXT NOT NULL DEFAULT '', created_at TEXT NOT NULL DEFAULT (datetime('now')) );", )?; info!(path, "Opened Sol state database"); Ok(Self { conn: Mutex::new(conn), }) } /// Open an in-memory database (for tests). pub fn open_memory() -> anyhow::Result { Self::open(":memory:") } // ========================================================================= // Conversations // ========================================================================= /// Get the conversation_id for a room, if one exists. pub fn get_conversation(&self, room_id: &str) -> Option<(String, u32)> { let conn = self.conn.lock().unwrap(); conn.query_row( "SELECT conversation_id, estimated_tokens FROM conversations WHERE room_id = ?1", params![room_id], |row| Ok((row.get::<_, String>(0)?, row.get::<_, u32>(1)?)), ) .ok() } /// Store or update a conversation mapping. pub fn upsert_conversation( &self, room_id: &str, conversation_id: &str, estimated_tokens: u32, ) { let conn = self.conn.lock().unwrap(); if let Err(e) = conn.execute( "INSERT INTO conversations (room_id, conversation_id, estimated_tokens) VALUES (?1, ?2, ?3) ON CONFLICT(room_id) DO UPDATE SET conversation_id = excluded.conversation_id, estimated_tokens = excluded.estimated_tokens", params![room_id, conversation_id, estimated_tokens], ) { warn!("Failed to upsert conversation: {e}"); } } /// Update token estimate for a conversation. pub fn update_tokens(&self, room_id: &str, estimated_tokens: u32) { let conn = self.conn.lock().unwrap(); let _ = conn.execute( "UPDATE conversations SET estimated_tokens = ?1 WHERE room_id = ?2", params![estimated_tokens, room_id], ); } /// Remove a conversation mapping (e.g., after compaction). pub fn delete_conversation(&self, room_id: &str) { let conn = self.conn.lock().unwrap(); let _ = conn.execute( "DELETE FROM conversations WHERE room_id = ?1", params![room_id], ); } /// Load all conversation mappings (for startup recovery). pub fn load_all_conversations(&self) -> Vec<(String, String, u32)> { let conn = self.conn.lock().unwrap(); let mut stmt = match conn.prepare( "SELECT room_id, conversation_id, estimated_tokens FROM conversations", ) { Ok(s) => s, Err(_) => return Vec::new(), }; stmt.query_map([], |row| { Ok(( row.get::<_, String>(0)?, row.get::<_, String>(1)?, row.get::<_, u32>(2)?, )) }) .ok() .map(|rows| rows.filter_map(|r| r.ok()).collect()) .unwrap_or_default() } // ========================================================================= // Agents // ========================================================================= /// Get the agent_id for a named agent. pub fn get_agent(&self, name: &str) -> Option { let conn = self.conn.lock().unwrap(); conn.query_row( "SELECT agent_id FROM agents WHERE name = ?1", params![name], |row| row.get(0), ) .ok() } /// Store or update an agent mapping. pub fn upsert_agent(&self, name: &str, agent_id: &str, model: &str) { let conn = self.conn.lock().unwrap(); if let Err(e) = conn.execute( "INSERT INTO agents (name, agent_id, model) VALUES (?1, ?2, ?3) ON CONFLICT(name) DO UPDATE SET agent_id = excluded.agent_id, model = excluded.model", params![name, agent_id, model], ) { warn!("Failed to upsert agent: {e}"); } } /// Remove an agent mapping. pub fn delete_agent(&self, name: &str) { let conn = self.conn.lock().unwrap(); let _ = conn.execute("DELETE FROM agents WHERE name = ?1", params![name]); } /// Load all agent mappings (for startup recovery). pub fn load_all_agents(&self) -> Vec<(String, String)> { let conn = self.conn.lock().unwrap(); let mut stmt = match conn.prepare("SELECT name, agent_id FROM agents") { Ok(s) => s, Err(_) => return Vec::new(), }; stmt.query_map([], |row| { Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)) }) .ok() .map(|rows| rows.filter_map(|r| r.ok()).collect()) .unwrap_or_default() } } #[cfg(test)] mod tests { use super::*; #[test] fn test_open_memory_db() { let store = Store::open_memory().unwrap(); assert!(store.load_all_conversations().is_empty()); assert!(store.load_all_agents().is_empty()); } #[test] fn test_conversation_crud() { let store = Store::open_memory().unwrap(); // Insert store.upsert_conversation("!room:x", "conv_abc", 100); let (conv_id, tokens) = store.get_conversation("!room:x").unwrap(); assert_eq!(conv_id, "conv_abc"); assert_eq!(tokens, 100); // Update tokens store.update_tokens("!room:x", 500); let (_, tokens) = store.get_conversation("!room:x").unwrap(); assert_eq!(tokens, 500); // Upsert (replace conversation_id) store.upsert_conversation("!room:x", "conv_def", 0); let (conv_id, tokens) = store.get_conversation("!room:x").unwrap(); assert_eq!(conv_id, "conv_def"); assert_eq!(tokens, 0); // Delete store.delete_conversation("!room:x"); assert!(store.get_conversation("!room:x").is_none()); } #[test] fn test_agent_crud() { let store = Store::open_memory().unwrap(); store.upsert_agent("sol-orchestrator", "ag_123", "mistral-medium-latest"); assert_eq!( store.get_agent("sol-orchestrator").unwrap(), "ag_123" ); // Update store.upsert_agent("sol-orchestrator", "ag_456", "mistral-medium-latest"); assert_eq!( store.get_agent("sol-orchestrator").unwrap(), "ag_456" ); // Delete store.delete_agent("sol-orchestrator"); assert!(store.get_agent("sol-orchestrator").is_none()); } #[test] fn test_load_all_conversations() { let store = Store::open_memory().unwrap(); store.upsert_conversation("!a:x", "conv_1", 10); store.upsert_conversation("!b:x", "conv_2", 20); store.upsert_conversation("!c:x", "conv_3", 30); let all = store.load_all_conversations(); assert_eq!(all.len(), 3); } #[test] fn test_load_all_agents() { let store = Store::open_memory().unwrap(); store.upsert_agent("orch", "ag_1", "medium"); store.upsert_agent("obs", "ag_2", "medium"); let all = store.load_all_agents(); assert_eq!(all.len(), 2); } #[test] fn test_nonexistent_keys_return_none() { let store = Store::open_memory().unwrap(); assert!(store.get_conversation("!nope:x").is_none()); assert!(store.get_agent("nope").is_none()); } }