magi-rs 0.3.0

//! This module provides a persistent memory system based on SQLite with encryption.

use crate::agent::messages::Message;
use crate::utils::crypto::{CryptoVault, ErrorCorrection, ReedSolomonCodec, SALT_LEN};
use anyhow::Result;
use async_trait::async_trait;
use rand::RngCore;
use rusqlite::{params, Connection, OptionalExtension};
use sha2::{Digest, Sha256};
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use zeroize::Zeroizing;

/// Length of the SHA-256 salt checksum stored alongside the salt (#15).
const SALT_CHECKSUM_LEN: usize = 4;

/// RS-encodes the per-DB salt together with a short SHA-256 checksum (#15). The
/// checksum lets the reader reject a salt that RS *mis-corrects* to a wrong but
/// valid-looking codeword (e.g. an all-zero block) — which RS error-correction
/// alone cannot detect — closing the #10 residual.
fn encode_salt(codec: &ReedSolomonCodec, salt: &[u8]) -> Vec<u8> {
    let mut payload = Vec::with_capacity(SALT_LEN + SALT_CHECKSUM_LEN);
    payload.extend_from_slice(salt);
    payload.extend_from_slice(&Sha256::digest(salt)[..SALT_CHECKSUM_LEN]);
    codec.encode(&payload)
}

/// Decodes and verifies a salt blob from [`encode_salt`]. Returns the 16-byte
/// salt only if RS-decode succeeds, the payload length is exact, and the checksum
/// matches; otherwise `None` (treated as absent → D6 self-heal).
fn decode_salt(codec: &ReedSolomonCodec, blob: &[u8]) -> Option<Vec<u8>> {
    let payload = codec.decode(blob, SALT_LEN + SALT_CHECKSUM_LEN).ok()?;
    if payload.len() != SALT_LEN + SALT_CHECKSUM_LEN {
        return None;
    }
    let (salt, checksum) = payload.split_at(SALT_LEN);
    let expected = Sha256::digest(salt);
    if expected[..SALT_CHECKSUM_LEN] == checksum[..] {
        Some(salt.to_vec())
    } else {
        None
    }
}

/// Trait defining the behavior of the agent's memory.
#[async_trait]
pub trait MemoryStore: Send + Sync {
    /// Creates a new session and returns its ID.
    async fn create_session(&self, project_name: &str) -> Result<String>;

    /// Adds a message to a specific session.
    async fn add_message(&self, session_id: &str, message: &Message) -> Result<()>;

    /// Retrieves all messages for a session.
    async fn get_messages(&self, session_id: &str) -> Result<Vec<Message>>;

    /// Lists all sessions.
    async fn list_sessions(&self) -> Result<Vec<(String, String)>>; // (id, project_name)

    /// Stores a persistent fact about the project.
    async fn set_knowledge(&self, key: &str, value: &str) -> Result<()>;

    /// Retrieves a persistent fact.
    async fn get_knowledge(&self, key: &str) -> Result<Option<String>>;

    /// Lists all known project keys.
    async fn list_knowledge_keys(&self) -> Result<Vec<String>>;
}

/// A persistent memory store using SQLite and CryptoVault for encryption.
pub struct EncryptedSqliteMemory {
    conn: Arc<Mutex<Connection>>,
    vault: CryptoVault,
    /// Data key derived **once** from the per-DB salt + master password (B′).
    ///
    /// `Zeroizing<Vec<u8>>` overwrites the key on drop. Derived in
    /// [`Self::new_with_vault`]; reused by every record so no per-record Argon2
    /// runs on the hot path.
    derived_key: Zeroizing<Vec<u8>>,
    /// `true` when construction discarded incompatible/corrupt on-disk content
    /// (the D6 reset fired on a non-empty DB). Surfaced to the user at startup (#11).
    was_reset: bool,
}

impl EncryptedSqliteMemory {
    /// Locks the connection, recovering the guard if the mutex was poisoned by a
    /// panic in another thread (the SQLite handle remains valid). Keeps
    /// persistence available instead of failing closed for the session (#8,
    /// supersedes the W11 error-on-poison behavior); the recovery is logged.
    fn locked_conn(&self) -> std::sync::MutexGuard<'_, Connection> {
        self.conn.lock().unwrap_or_else(|poisoned| {
            use std::sync::atomic::{AtomicBool, Ordering};
            // Warn once per process: a persistently-poisoned mutex would otherwise
            // spam stderr on every op (and disrupt the TUI alternate screen).
            static POISON_WARNED: AtomicBool = AtomicBool::new(false);
            if !POISON_WARNED.swap(true, Ordering::Relaxed) {
                eprintln!(
                    "WARNING: database connection mutex was poisoned by a panic in another \
                     thread; recovering the connection and continuing (further occurrences \
                     suppressed)."
                );
            }
            poisoned.into_inner()
        })
    }

    /// Collects raw `(role, blob)` rows for a session under the connection lock.
    ///
    /// The lock is held only for the duration of the SELECT and the iterator
    /// drain; it is released before any decryption happens (audit finding W12).
    fn collect_message_rows(&self, session_id: &str) -> Result<Vec<(String, String)>> {
        let conn = self.locked_conn();
        let mut stmt = conn.prepare(
            "SELECT role, content_blob FROM messages WHERE session_id = ? ORDER BY created_at ASC",
        )?;
        let mapped = stmt.query_map(params![session_id], |row| {
            Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
        })?;
        let mut collected = Vec::new();
        for row in mapped {
            collected.push(row?);
        }
        Ok(collected)
    }

    /// Decrypts pre-collected `(role, blob)` rows into [`Message`]s.
    ///
    /// Holds **no** database lock: callers must collect rows and release the
    /// connection guard before invoking this, so per-row Argon2 key derivation
    /// never serializes other DB callers (audit finding W12).
    fn decrypt_rows(&self, rows: Vec<(String, String)>) -> Result<Vec<Message>> {
        let mut messages = Vec::with_capacity(rows.len());
        for (role_str, blob) in rows {
            let decrypted = self
                .vault
                .decrypt_with_key(&self.derived_key, &blob)
                .map_err(|e| anyhow::anyhow!("Decryption failed: {}", e))?;
            let content = serde_json::from_str(&decrypted)?;
            let role = match role_str.as_str() {
                "User" => crate::agent::messages::Role::User,
                _ => crate::agent::messages::Role::Assistant,
            };
            messages.push(Message { role, content });
        }
        Ok(messages)
    }

    pub fn new(path: PathBuf, master_password: String) -> Result<Self> {
        Self::new_with_vault(path, master_password, CryptoVault::default())
    }

    /// Whether the on-disk DB had real content discarded (reset to fresh) during
    /// construction (#11). The caller surfaces this to the user at startup.
    pub fn was_reset(&self) -> bool {
        self.was_reset
    }

    /// Constructor that accepts a custom [`CryptoVault`] (e.g. a counting KDF in
    /// tests). Derives the data key **once** from the per-DB salt and caches it.
    pub(crate) fn new_with_vault(
        path: PathBuf,
        master_password: String,
        vault: CryptoVault,
    ) -> Result<Self> {
        let mut conn = Connection::open(path)?;

        // MAGI FIX: Enable WAL mode for high concurrency
        // We use query_row because execute fails for pragmas that return values in some drivers
        let _: String = conn.query_row("PRAGMA journal_mode = WAL", [], |row| row.get(0))?;
        conn.execute("PRAGMA synchronous = NORMAL", [])?;
        // Set a busy timeout to prevent "database is locked" errors during contention
        conn.busy_timeout(std::time::Duration::from_secs(5))?;

        conn.execute(
            "CREATE TABLE IF NOT EXISTS sessions (
                id TEXT PRIMARY KEY,
                project_name TEXT NOT NULL,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )",
            [],
        )?;
        conn.execute(
            "CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                role TEXT NOT NULL,
                content_blob TEXT NOT NULL,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY(session_id) REFERENCES sessions(id)
            )",
            [],
        )?;
        conn.execute(
            "CREATE TABLE IF NOT EXISTS knowledge (
                key TEXT PRIMARY KEY,
                value_blob TEXT NOT NULL,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )",
            [],
        )?;
        conn.execute(
            "CREATE TABLE IF NOT EXISTS vault_meta (
                key TEXT PRIMARY KEY,
                value BLOB NOT NULL
            )",
            [],
        )?;

        // B′: one per-DB salt → one key derivation. The salt is **RS-encoded with a
        // SHA-256 checksum** on disk (#10 + #15): minor bit-rot is corrected on read,
        // and a salt that is absent, fails to RS-decode, or fails the checksum (incl.
        // a valid-codeword mis-correction such as an all-zero block) is treated as
        // **absent**, so D6 self-heals (reset + re-bootstrap) instead of bricking.
        // Per D6 the reset wipes content in a single transaction (atomic) and warns
        // only when real (non-empty) history is discarded, so it is observable.
        let salt_codec = ReedSolomonCodec::default();
        let valid_salt: Option<Vec<u8>> = conn
            .query_row("SELECT value FROM vault_meta WHERE key = 'salt'", [], |r| {
                r.get::<_, Vec<u8>>(0)
            })
            .optional()?
            // `decode_salt` returns the salt only if RS-decode succeeds, the payload
            // length is exact, AND the SHA-256 checksum matches (#15); anything else
            // (absent, corrupt, a non-RS raw value from a pre-#10 DB, or a
            // checksum-failing mis-correction) maps to `None` to trigger the D6
            // self-heal below rather than deriving a wrong key from a bad salt.
            .and_then(|blob| decode_salt(&salt_codec, &blob));
        let mut was_reset = false;
        let salt: Vec<u8> = match valid_salt {
            Some(s) => s,
            None => {
                let had_rows: i64 = conn.query_row(
                    "SELECT (SELECT COUNT(*) FROM sessions) \
                     + (SELECT COUNT(*) FROM messages) \
                     + (SELECT COUNT(*) FROM knowledge)",
                    [],
                    |r| r.get(0),
                )?;

                let mut new_salt = vec![0u8; SALT_LEN];
                rand::rngs::OsRng.fill_bytes(&mut new_salt);
                let salt_blob = encode_salt(&salt_codec, &new_salt);

                // IMMEDIATE acquires the write lock at BEGIN (#12): a process racing
                // the same brand-new DB blocks here until the winner commits, then
                // re-reads and ADOPTS the winner's salt — instead of failing with
                // SQLITE_BUSY_SNAPSHOT, which a DEFERRED tx would after its read.
                let tx =
                    conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
                // Re-check under that write lock: a racing opener may have written a
                // VALID salt between our pre-tx read and here. If so, adopt it (every
                // opener converges on the same salt) and leave content untouched. Only
                // when no valid salt exists (absent, or present-but-corrupt per the #10
                // self-heal) do we wipe incompatible content and bootstrap our own.
                let current: Option<Vec<u8>> = tx
                    .query_row("SELECT value FROM vault_meta WHERE key = 'salt'", [], |r| {
                        r.get(0)
                    })
                    .optional()?;
                let current_valid = current.and_then(|b| decode_salt(&salt_codec, &b));
                let (salt, wiped) = match current_valid {
                    Some(existing) => (existing, false),
                    None => {
                        tx.execute("DELETE FROM messages", [])?;
                        tx.execute("DELETE FROM knowledge", [])?;
                        tx.execute("DELETE FROM sessions", [])?;
                        // OR REPLACE: in the self-heal case the (corrupt) row exists.
                        tx.execute(
                            "INSERT OR REPLACE INTO vault_meta (key, value) VALUES ('salt', ?1)",
                            params![salt_blob],
                        )?;
                        (new_salt, true)
                    }
                };
                tx.commit()?;

                was_reset = wiped && had_rows > 0;
                if was_reset {
                    eprintln!(
                        "WARNING: existing on-disk history used an incompatible or corrupt \
                         encryption salt and has been reset (fresh start). This is expected \
                         after upgrading the storage format."
                    );
                }
                salt
            }
        };

        // Derive the data key exactly once; the incoming password is zeroized
        // after derivation.
        let password = Zeroizing::new(master_password);
        let derived_key = vault
            .derive_key(&password, &salt)
            .map_err(|e| anyhow::anyhow!("Key derivation failed: {}", e))?;

        Ok(Self {
            conn: Arc::new(Mutex::new(conn)),
            vault,
            derived_key,
            was_reset,
        })
    }
}

#[async_trait]
impl MemoryStore for EncryptedSqliteMemory {
    async fn create_session(&self, project_name: &str) -> Result<String> {
        let id = uuid::Uuid::new_v4().to_string();
        let conn = self.locked_conn();
        conn.execute(
            "INSERT INTO sessions (id, project_name) VALUES (?1, ?2)",
            params![id, project_name],
        )?;
        Ok(id)
    }

    async fn add_message(&self, session_id: &str, message: &Message) -> Result<()> {
        let json_content = serde_json::to_string(&message.content)?;
        let encrypted = self
            .vault
            .encrypt_with_key(&self.derived_key, &json_content)
            .map_err(|e| anyhow::anyhow!("Encryption failed: {}", e))?;

        let conn = self.locked_conn();
        conn.execute(
            "INSERT INTO messages (session_id, role, content_blob) VALUES (?1, ?2, ?3)",
            params![session_id, format!("{:?}", message.role), encrypted],
        )?;
        Ok(())
    }

    async fn get_messages(&self, session_id: &str) -> Result<Vec<Message>> {
        let raw_rows = self.collect_message_rows(session_id)?;
        self.decrypt_rows(raw_rows)
    }

    async fn list_sessions(&self) -> Result<Vec<(String, String)>> {
        let conn = self.locked_conn();
        let mut stmt =
            conn.prepare("SELECT id, project_name FROM sessions ORDER BY created_at DESC")?;
        let rows = stmt.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?;

        let mut sessions = Vec::new();
        for row in rows {
            sessions.push(row?);
        }
        Ok(sessions)
    }

    async fn set_knowledge(&self, key: &str, value: &str) -> Result<()> {
        let encrypted = self
            .vault
            .encrypt_with_key(&self.derived_key, value)
            .map_err(|e| anyhow::anyhow!("Encryption failed: {}", e))?;

        let conn = self.locked_conn();
        conn.execute(
            "INSERT OR REPLACE INTO knowledge (key, value_blob, updated_at) VALUES (?1, ?2, CURRENT_TIMESTAMP)",
            params![key, encrypted],
        )?;
        Ok(())
    }

    async fn get_knowledge(&self, key: &str) -> Result<Option<String>> {
        let conn = self.locked_conn();
        let mut stmt = conn.prepare("SELECT value_blob FROM knowledge WHERE key = ?")?;

        let res = stmt.query_row(params![key], |row| row.get::<_, String>(0));

        match res {
            Ok(blob) => {
                let decrypted = self
                    .vault
                    .decrypt_with_key(&self.derived_key, &blob)
                    .map_err(|e| anyhow::anyhow!("Decryption failed: {}", e))?;
                Ok(Some(decrypted))
            }
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(e) => Err(anyhow::anyhow!("Database error: {}", e)),
        }
    }

    async fn list_knowledge_keys(&self) -> Result<Vec<String>> {
        let conn = self.locked_conn();
        let mut stmt = conn.prepare("SELECT key FROM knowledge ORDER BY key ASC")?;
        let rows = stmt.query_map([], |row| row.get(0))?;

        let mut keys = Vec::new();
        for row in rows {
            keys.push(row?);
        }
        Ok(keys)
    }
}

#[cfg(test)]
impl EncryptedSqliteMemory {
    pub(crate) fn conn_for_test(&self) -> &Arc<Mutex<Connection>> {
        &self.conn
    }

    pub(crate) fn collect_message_rows_for_test(
        &self,
        session_id: &str,
    ) -> Result<Vec<(String, String)>> {
        self.collect_message_rows(session_id)
    }

    pub(crate) fn derived_key_type_for_test(&self) -> &zeroize::Zeroizing<Vec<u8>> {
        &self.derived_key
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::utils::crypto::{
        Aes256GcmSivCipher, Argon2Kdf, CryptoError, CryptoVault, KeyDerivation, ReedSolomonCodec,
    };
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Arc;
    use tempfile::NamedTempFile;

    /// KDF that counts derivations and delegates to the real Argon2id.
    struct CountingKdf {
        inner: Argon2Kdf,
        calls: Arc<AtomicUsize>,
    }
    impl KeyDerivation for CountingKdf {
        fn derive_key(
            &self,
            password: &[u8],
            salt: &[u8],
            output_len: usize,
        ) -> std::result::Result<zeroize::Zeroizing<Vec<u8>>, CryptoError> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            self.inner.derive_key(password, salt, output_len)
        }
    }

    #[tokio::test]
    async fn test_key_is_derived_exactly_once_for_session_load() {
        // S-6 (load-bearing): construct + N adds + get_messages => 1 Argon2 call.
        let tmp = NamedTempFile::new().unwrap();
        let calls = Arc::new(AtomicUsize::new(0));
        let vault = CryptoVault::new(
            Box::new(CountingKdf {
                inner: Argon2Kdf,
                calls: calls.clone(),
            }),
            Box::new(Aes256GcmSivCipher),
            Box::new(ReedSolomonCodec::default()),
        );
        let memory = EncryptedSqliteMemory::new_with_vault(
            tmp.path().to_path_buf(),
            "pw".to_string(),
            vault,
        )
        .unwrap();
        let sid = memory.create_session("p").await.unwrap();
        for i in 0..5 {
            memory
                .add_message(&sid, &Message::user(&format!("m{i}")))
                .await
                .unwrap();
        }
        let msgs = memory.get_messages(&sid).await.unwrap();
        assert_eq!(msgs.len(), 5);
        assert_eq!(
            calls.load(Ordering::SeqCst),
            1,
            "Argon2 must run exactly once (B′), not per record"
        );
    }

    #[tokio::test]
    async fn test_was_reset_flag_reflects_content_discard() {
        // S-1 (#11): a legacy DB (rows, no salt) that gets reset reports
        // was_reset() == true; a fresh DB reports false.
        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();
        {
            let conn = Connection::open(&path).unwrap();
            conn.execute(
                "CREATE TABLE sessions (id TEXT PRIMARY KEY, project_name TEXT NOT NULL, \
                 created_at DATETIME DEFAULT CURRENT_TIMESTAMP)",
                [],
            )
            .unwrap();
            conn.execute(
                "CREATE TABLE messages (id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL, \
                 role TEXT NOT NULL, content_blob TEXT NOT NULL, created_at DATETIME DEFAULT CURRENT_TIMESTAMP)",
                [],
            )
            .unwrap();
            conn.execute(
                "INSERT INTO sessions (id, project_name) VALUES ('old', 'legacy')",
                [],
            )
            .unwrap();
            conn.execute(
                "INSERT INTO messages (session_id, role, content_blob) VALUES ('old', 'User', 'X')",
                [],
            )
            .unwrap();
        }
        let legacy = EncryptedSqliteMemory::new(path, "pw".to_string()).unwrap();
        assert!(
            legacy.was_reset(),
            "a legacy DB that discarded content must report was_reset()"
        );

        let tmp2 = NamedTempFile::new().unwrap();
        let fresh =
            EncryptedSqliteMemory::new(tmp2.path().to_path_buf(), "pw".to_string()).unwrap();
        assert!(!fresh.was_reset(), "a fresh DB must not report was_reset()");
    }

    #[tokio::test]
    async fn test_legacy_db_without_salt_is_reset_on_open() {
        // S-7: a pre-B′ DB (rows present, no vault_meta salt) is wiped on open.
        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();
        {
            let conn = Connection::open(&path).unwrap();
            conn.execute(
                "CREATE TABLE sessions (id TEXT PRIMARY KEY, project_name TEXT NOT NULL, \
                 created_at DATETIME DEFAULT CURRENT_TIMESTAMP)",
                [],
            )
            .unwrap();
            conn.execute(
                "CREATE TABLE messages (id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL, \
                 role TEXT NOT NULL, content_blob TEXT NOT NULL, created_at DATETIME DEFAULT CURRENT_TIMESTAMP)",
                [],
            )
            .unwrap();
            conn.execute(
                "INSERT INTO sessions (id, project_name) VALUES ('old', 'legacy')",
                [],
            )
            .unwrap();
            conn.execute(
                "INSERT INTO messages (session_id, role, content_blob) VALUES ('old', 'User', 'OLD_BLOB')",
                [],
            )
            .unwrap();
        }

        let memory = EncryptedSqliteMemory::new(path, "pw".to_string()).unwrap();
        assert!(
            memory.list_sessions().await.unwrap().is_empty(),
            "legacy rows must be wiped on open (D6 fresh-start)"
        );
        let sid = memory.create_session("fresh").await.unwrap();
        memory
            .add_message(&sid, &Message::user("new"))
            .await
            .unwrap();
        assert_eq!(memory.get_messages(&sid).await.unwrap().len(), 1);
    }

    #[tokio::test]
    async fn test_salt_persists_across_reopen_same_password_roundtrips() {
        // S-8: salt persists => same password round-trips; different password fails.
        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();
        let sid;
        {
            let memory = EncryptedSqliteMemory::new(path.clone(), "P".to_string()).unwrap();
            sid = memory.create_session("p").await.unwrap();
            memory
                .add_message(&sid, &Message::user("persisted"))
                .await
                .unwrap();
        }
        {
            let memory = EncryptedSqliteMemory::new(path.clone(), "P".to_string()).unwrap();
            assert_eq!(
                memory.get_messages(&sid).await.unwrap(),
                vec![Message::user("persisted")]
            );
        }
        {
            let memory = EncryptedSqliteMemory::new(path, "P-different".to_string()).unwrap();
            let res = memory.get_messages(&sid).await;
            assert!(res.is_err());
            assert!(res.unwrap_err().to_string().contains("Decryption failed"));
        }
    }

    #[tokio::test]
    async fn test_corrupt_salt_triggers_self_heal_reset() {
        // S-2: an invalid/unrecoverable salt (here: a raw, non-RS-encoded value)
        // is treated as absent so D6 self-heals (reset) instead of bricking the DB.
        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();
        {
            let memory = EncryptedSqliteMemory::new(path.clone(), "P".to_string()).unwrap();
            let sid = memory.create_session("p").await.unwrap();
            memory
                .add_message(&sid, &Message::user("orig"))
                .await
                .unwrap();
        }
        // Overwrite the salt with a raw, non-RS-encoded value (irrecoverable).
        {
            let conn = Connection::open(&path).unwrap();
            conn.execute(
                "UPDATE vault_meta SET value = ?1 WHERE key = 'salt'",
                params![vec![0x42u8; SALT_LEN]],
            )
            .unwrap();
        }
        // Reopen: invalid salt -> treated absent -> D6 reset, no error.
        let memory = EncryptedSqliteMemory::new(path, "P".to_string()).unwrap();
        assert!(
            memory.list_sessions().await.unwrap().is_empty(),
            "an invalid/unrecoverable salt must self-heal via a D6 reset"
        );
        let fresh = memory.create_session("fresh").await.unwrap();
        memory
            .add_message(&fresh, &Message::user("new"))
            .await
            .unwrap();
        assert_eq!(memory.get_messages(&fresh).await.unwrap().len(), 1);
    }

    #[tokio::test]
    async fn test_various_invalid_salt_shapes_self_heal() {
        // Strengthens S-2: every salt shape that RS-decode REJECTS (empty,
        // raw/non-RS-encoded, truncated below one block) must self-heal via a D6
        // reset, never validate to a wrong salt. A valid-codeword corruption such
        // as an all-zero block is covered by `test_all_zero_salt_self_heals_via_checksum`
        // (#15, now closed via the salt checksum).
        for bad in [Vec::<u8>::new(), vec![0x42u8; SALT_LEN], vec![0x07u8; 30]] {
            let tmp = NamedTempFile::new().unwrap();
            let path = tmp.path().to_path_buf();
            {
                let m = EncryptedSqliteMemory::new(path.clone(), "P".to_string()).unwrap();
                let s = m.create_session("p").await.unwrap();
                m.add_message(&s, &Message::user("x")).await.unwrap();
            }
            {
                let conn = Connection::open(&path).unwrap();
                conn.execute(
                    "UPDATE vault_meta SET value = ?1 WHERE key = 'salt'",
                    params![bad],
                )
                .unwrap();
            }
            let m = EncryptedSqliteMemory::new(path, "P".to_string()).unwrap();
            assert!(
                m.list_sessions().await.unwrap().is_empty(),
                "an RS-rejected invalid salt must self-heal via a D6 reset"
            );
        }
    }

    #[tokio::test]
    async fn test_all_zero_salt_self_heals_via_checksum() {
        // B-S1 (#15): an all-zero salt blob is a VALID RS codeword that decodes to a
        // zero payload (the #10 mis-correction residual — RS alone would "accept"
        // it). The salt checksum won't match sha256(zero-salt) → treated as invalid
        // → D6 self-heal, instead of adopting a wrong (zero) salt and bricking.
        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();
        {
            let m = EncryptedSqliteMemory::new(path.clone(), "P".to_string()).unwrap();
            let s = m.create_session("p").await.unwrap();
            m.add_message(&s, &Message::user("orig")).await.unwrap();
        }
        {
            // A salt blob is RS(salt[16] ‖ checksum[4]) = 20 data + 32 parity = 52
            // bytes; an all-zero blob of that size is a valid (zero) codeword.
            let conn = Connection::open(&path).unwrap();
            conn.execute(
                "UPDATE vault_meta SET value = ?1 WHERE key = 'salt'",
                params![vec![0u8; 52]],
            )
            .unwrap();
        }
        let m = EncryptedSqliteMemory::new(path, "P".to_string()).unwrap();
        assert!(
            m.list_sessions().await.unwrap().is_empty(),
            "an all-zero (checksum-mismatched) salt must self-heal, not be adopted"
        );
        let s = m.create_session("fresh").await.unwrap();
        m.add_message(&s, &Message::user("new")).await.unwrap();
        assert_eq!(m.get_messages(&s).await.unwrap().len(), 1);
    }

    #[tokio::test]
    async fn test_minor_salt_bitrot_is_corrected_and_history_survives() {
        // S-3: an RS-encoded salt with a few flipped bytes is corrected on read,
        // so the derived key is unchanged and prior history still decrypts.
        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();
        let sid;
        {
            let memory = EncryptedSqliteMemory::new(path.clone(), "P".to_string()).unwrap();
            sid = memory.create_session("p").await.unwrap();
            memory
                .add_message(&sid, &Message::user("survives"))
                .await
                .unwrap();
        }
        // Flip a few bytes of the stored salt blob (within RS correction capacity).
        {
            let conn = Connection::open(&path).unwrap();
            let mut blob: Vec<u8> = conn
                .query_row("SELECT value FROM vault_meta WHERE key = 'salt'", [], |r| {
                    r.get(0)
                })
                .unwrap();
            for b in blob.iter_mut().take(3) {
                *b ^= 0xAA;
            }
            conn.execute(
                "UPDATE vault_meta SET value = ?1 WHERE key = 'salt'",
                params![blob],
            )
            .unwrap();
        }
        // Reopen: RS corrects the salt -> same key -> history survives.
        let memory = EncryptedSqliteMemory::new(path, "P".to_string()).unwrap();
        assert_eq!(
            memory.get_messages(&sid).await.unwrap(),
            vec![Message::user("survives")],
            "RS-encoded salt must self-correct minor bit-rot, preserving history"
        );
    }

    #[tokio::test]
    async fn test_encrypted_sqlite_memory() {
        let tmp_file = NamedTempFile::new().unwrap();
        let path = tmp_file.path().to_path_buf();
        let password = "master_key_123";

        let memory = EncryptedSqliteMemory::new(path, password.to_string()).unwrap();
        let sid = memory.create_session("test_proj").await.unwrap();

        let msg = Message::user("Hello secure world");
        memory.add_message(&sid, &msg).await.unwrap();

        let msgs = memory.get_messages(&sid).await.unwrap();
        assert_eq!(msgs.len(), 1);
        assert_eq!(msgs[0], msg);

        // Verify encryption (raw read)
        let conn = Connection::open(tmp_file.path()).unwrap();
        let blob: String = conn
            .query_row("SELECT content_blob FROM messages LIMIT 1", [], |r| {
                r.get(0)
            })
            .unwrap();
        assert!(
            !blob.contains("Hello"),
            "Database should contain encrypted blob, not plaintext"
        );

        // Verify list_sessions (to clear dead code warning)
        let sessions = memory.list_sessions().await.unwrap();
        assert_eq!(sessions.len(), 1);
        assert_eq!(sessions[0].1, "test_proj");
    }

    #[tokio::test]
    async fn test_project_knowledge_persistence() {
        let tmp_file = NamedTempFile::new().unwrap();
        let path = tmp_file.path().to_path_buf();
        let password = "knowledge_key_123".to_string();

        let memory = EncryptedSqliteMemory::new(path, password).unwrap();

        memory
            .set_knowledge("architecture", "Clean hex with encrypted SQLite")
            .await
            .unwrap();

        let fact = memory.get_knowledge("architecture").await.unwrap();
        assert_eq!(fact.unwrap(), "Clean hex with encrypted SQLite");

        // Verify multiple keys
        memory.set_knowledge("port", "54545").await.unwrap();
        let keys = memory.list_knowledge_keys().await.unwrap();
        assert_eq!(keys.len(), 2);
        assert!(keys.contains(&"architecture".to_string()));
        assert!(keys.contains(&"port".to_string()));
    }

    #[tokio::test]
    async fn test_sqlite_concurrency_stress() {
        let tmp_file = tempfile::NamedTempFile::new().unwrap();
        let path = tmp_file.path().to_path_buf();
        let memory = Arc::new(EncryptedSqliteMemory::new(path, "stress_pass".to_string()).unwrap());

        let mut handles = vec![];
        for i in 0..20 {
            let mem_clone = memory.clone();
            handles.push(tokio::spawn(async move {
                let key = format!("key_{}", i);
                let val = format!("val_{}", i);
                mem_clone.set_knowledge(&key, &val).await
            }));
        }

        for h in handles {
            let res = h.await.unwrap();
            assert!(res.is_ok(), "Concurrent write failed: {:?}", res.err());
        }

        let keys = memory.list_knowledge_keys().await.unwrap();
        assert_eq!(keys.len(), 20);
    }

    #[tokio::test]
    async fn test_poisoned_lock_recovers_and_continues() {
        // A-S1 (#8, supersedes W11): a poisoned mutex is recovered (into_inner) so
        // persistence keeps working instead of failing closed for the session.
        let tmp_file = NamedTempFile::new().unwrap();
        let path = tmp_file.path().to_path_buf();
        let memory = EncryptedSqliteMemory::new(path, "pw".to_string()).unwrap();

        let conn = memory.conn_for_test().clone();
        let _ = std::thread::spawn(move || {
            let _guard = conn.lock().unwrap();
            panic!("intentional poison");
        })
        .join();

        // The lock is now poisoned; operations must recover and succeed.
        assert!(
            memory.list_sessions().await.is_ok(),
            "a poisoned lock must be recovered, not fail closed"
        );
        let sid = memory.create_session("after-poison").await.unwrap();
        assert!(!sid.is_empty());
        assert_eq!(
            memory.list_sessions().await.unwrap().len(),
            1,
            "persistence continues working after lock recovery"
        );
    }

    #[tokio::test]
    async fn test_derived_key_matches_persisted_salt() {
        // B-S1 (#12): the cached key is always derived from the salt actually
        // stored on disk, so a concurrent-bootstrap winner can never diverge.
        let tmp = NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();
        let memory = EncryptedSqliteMemory::new(path.clone(), "P".to_string()).unwrap();

        let blob: Vec<u8> = {
            let conn = Connection::open(&path).unwrap();
            conn.query_row("SELECT value FROM vault_meta WHERE key = 'salt'", [], |r| {
                r.get(0)
            })
            .unwrap()
        };
        let codec = ReedSolomonCodec::default();
        let salt = decode_salt(&codec, &blob).expect("persisted salt must decode + verify");
        let expected = CryptoVault::default().derive_key("P", &salt).unwrap();
        assert_eq!(
            memory.derived_key_type_for_test().as_slice(),
            expected.as_slice(),
            "cached derived_key must be derived from the persisted salt"
        );
    }

    #[test]
    fn test_decode_salt_rejects_checksum_mismatch_and_roundtrips() {
        // #15 helper unit test: a valid RS codeword whose payload has a WRONG
        // checksum is rejected (the mis-correction guard), even though RS-decode
        // itself succeeds; and a correctly-encoded salt round-trips.
        let codec = ReedSolomonCodec::default();
        let mut bad_payload = vec![7u8; SALT_LEN]; // non-zero salt
        bad_payload.extend_from_slice(&[0u8; SALT_CHECKSUM_LEN]); // deliberately wrong checksum
        let bad_blob = codec.encode(&bad_payload);
        assert!(
            decode_salt(&codec, &bad_blob).is_none(),
            "a valid codeword with a mismatched checksum must be rejected"
        );

        let salt = vec![7u8; SALT_LEN];
        assert_eq!(
            decode_salt(&codec, &encode_salt(&codec, &salt)).unwrap(),
            salt,
            "a correctly-encoded salt must round-trip"
        );
    }

    #[tokio::test]
    async fn test_get_messages_does_not_hold_lock_during_decrypt() {
        let tmp_file = NamedTempFile::new().unwrap();
        let path = tmp_file.path().to_path_buf();
        let memory = Arc::new(EncryptedSqliteMemory::new(path, "pw".to_string()).unwrap());
        let sid = memory.create_session("p").await.unwrap();

        for i in 0..4 {
            memory
                .add_message(&sid, &Message::user(&format!("message number {i}")))
                .await
                .unwrap();
        }

        let reader = {
            let m = memory.clone();
            let s = sid.clone();
            tokio::spawn(async move { m.get_messages(&s).await })
        };
        let writer = {
            let m = memory.clone();
            tokio::spawn(async move { m.create_session("concurrent").await })
        };

        let msgs = reader.await.unwrap().unwrap();
        let new_sid = writer.await.unwrap().unwrap();

        assert_eq!(
            msgs.len(),
            4,
            "all messages decrypt correctly after lock-drop refactor"
        );
        assert!(
            !new_sid.is_empty(),
            "a concurrent write completes; lock is not held across decrypt"
        );
        assert_eq!(msgs[0], Message::user("message number 0"));
    }

    #[tokio::test]
    async fn test_decrypt_rows_runs_without_connection_lock() {
        let tmp_file = NamedTempFile::new().unwrap();
        let memory =
            EncryptedSqliteMemory::new(tmp_file.path().to_path_buf(), "pw".to_string()).unwrap();
        let sid = memory.create_session("p").await.unwrap();
        memory
            .add_message(&sid, &Message::user("hi"))
            .await
            .unwrap();

        let raw = memory.collect_message_rows_for_test(&sid).unwrap();
        let msgs = memory.decrypt_rows(raw).unwrap();
        assert_eq!(msgs, vec![Message::user("hi")]);
    }

    #[tokio::test]
    async fn test_derived_key_field_is_zeroizing_and_roundtrips() {
        let tmp_file = NamedTempFile::new().unwrap();
        let path = tmp_file.path().to_path_buf();

        let memory = EncryptedSqliteMemory::new(path, "zeroizing_pw".to_string()).unwrap();
        let sid = memory.create_session("p").await.unwrap();
        memory
            .add_message(&sid, &Message::user("secret payload"))
            .await
            .unwrap();

        let _assert_type: &zeroize::Zeroizing<Vec<u8>> = memory.derived_key_type_for_test();

        let msgs = memory.get_messages(&sid).await.unwrap();
        assert_eq!(msgs, vec![Message::user("secret payload")]);
    }
}