agentdb/

db.rs

use crate::audit::AuditStore;
use crate::context::ContextStore;
use crate::conversations::ConversationStore;
use crate::error::Result;
use crate::fts::FullTextStore;
use crate::hybrid::{HybridQuery, HybridResult, HybridStore, TriModalQuery, TriModalResult};
use crate::labels::LabelStore;
use crate::memory::MemoryGraph;
use crate::prompts::PromptStore;
use crate::schema;
use crate::tools::ToolStore;
use crate::traces::TraceStore;
use crate::vectors::VectorStore;
use crate::workflows::WorkflowStore;
use rusqlite::Connection;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

/// The main AgentDB connection — your single-file AI database.
#[derive(Clone)]
pub struct AgentDB {
    conn: Arc<Mutex<Connection>>,
    closed: Arc<AtomicBool>,
}

impl AgentDB {
    /// Open or create an AgentDB database. Use `":memory:"` for tests.
    pub fn open(path: &str) -> Result<Self> {
        let conn = Connection::open(path)?;
        schema::bootstrap(&conn)?;
        schema::check_version(&conn)?;
        Ok(Self {
            conn: Arc::new(Mutex::new(conn)),
            closed: Arc::new(AtomicBool::new(false)),
        })
    }

    /// Access the vector store layer
    pub fn vectors(&self) -> VectorStore {
        VectorStore::new(Arc::clone(&self.conn))
    }

    /// Access the memory graph layer
    pub fn memory(&self) -> MemoryGraph {
        MemoryGraph::new(Arc::clone(&self.conn))
    }

    /// Access the full-text search layer
    pub fn fts(&self) -> FullTextStore {
        FullTextStore::new(Arc::clone(&self.conn))
    }

    /// Access the conversation / message-threading layer
    pub fn conversations(&self) -> ConversationStore {
        ConversationStore::new(Arc::clone(&self.conn))
    }

    /// Access the workflow persistence layer
    pub fn workflows(&self) -> WorkflowStore {
        WorkflowStore::new(Arc::clone(&self.conn))
    }

    /// Access the reasoning-trace layer
    pub fn traces(&self) -> TraceStore {
        TraceStore::new(Arc::clone(&self.conn))
    }

    /// Access the tool registry and call log
    pub fn tools(&self) -> ToolStore {
        ToolStore::new(Arc::clone(&self.conn))
    }

    /// Access the immutable audit log
    pub fn audit(&self) -> AuditStore {
        AuditStore::new(Arc::clone(&self.conn))
    }

    /// Access the token-budgeted context window manager
    pub fn context(&self) -> ContextStore {
        ContextStore::new(Arc::clone(&self.conn))
    }

    /// Access the versioned prompt template store
    pub fn prompts(&self) -> PromptStore {
        PromptStore::new(Arc::clone(&self.conn))
    }

    /// Access the data classification / privacy label store
    pub fn labels(&self) -> LabelStore {
        LabelStore::new(Arc::clone(&self.conn))
    }

    /// Run a tri-modal graph + vector + FTS query
    pub fn tri_modal_query(&self, q: &TriModalQuery) -> Result<Vec<TriModalResult>> {
        let dim: usize = {
            let conn = self.conn.lock().unwrap();
            conn.query_row(
                "SELECT dim FROM _adb_collections WHERE name = ?1",
                rusqlite::params![q.collection],
                |r| r.get::<_, i64>(0).map(|v| v as usize),
            )
            .unwrap_or(q.embedding.len())
        };
        let col = self.vectors().collection(&q.collection, dim)?;
        let store = HybridStore::new(Arc::clone(&self.conn));
        store.tri_modal_query(q, &col)
    }

    /// Run a hybrid graph + vector query
    pub fn hybrid_query(&self, q: HybridQuery) -> Result<Vec<HybridResult>> {
        let dim: usize = {
            let conn = self.conn.lock().unwrap();
            conn.query_row(
                "SELECT dim FROM _adb_collections WHERE name = ?1",
                rusqlite::params![q.collection],
                |r| r.get::<_, i64>(0).map(|v| v as usize),
            )
            .unwrap_or(q.embedding.len())
        };
        let col = self.vectors().collection(q.collection, dim)?;
        let store = HybridStore::new(Arc::clone(&self.conn));
        store.query(q, &col)
    }

    /// Execute a raw SQL statement
    pub fn execute(&self, sql: &str) -> Result<usize> {
        let conn = self.conn.lock().unwrap();
        Ok(conn.execute(sql, [])?)
    }

    /// Execute a parameterized SQL statement
    pub fn execute_params(&self, sql: &str, params: &[&dyn rusqlite::ToSql]) -> Result<usize> {
        let conn = self.conn.lock().unwrap();
        Ok(conn.execute(sql, params)?)
    }

    /// Run multiple operations atomically inside a single SQLite transaction.
    ///
    /// The closure receives a [`rusqlite::Transaction`] and may perform any
    /// number of reads or writes.  If the closure returns `Ok`, the transaction
    /// is committed; if it returns `Err` (or panics), the transaction is rolled
    /// back automatically.
    ///
    /// # Example
    /// ```rust,no_run
    /// # use agentdb::AgentDB;
    /// let db = AgentDB::open(":memory:").unwrap();
    /// db.transaction(|tx| {
    ///     tx.execute("INSERT INTO _adb_nodes (id, kind, data, created_at, updated_at) VALUES ('x','tag','{}',0,0)", [])?;
    ///     tx.execute("INSERT INTO _adb_nodes (id, kind, data, created_at, updated_at) VALUES ('y','tag','{}',0,0)", [])?;
    ///     Ok(())
    /// }).unwrap();
    /// ```
    pub fn transaction<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&rusqlite::Transaction) -> Result<T>,
    {
        let mut conn = self.conn.lock().unwrap();
        let tx = conn.transaction()?;
        let result = f(&tx)?;
        tx.commit()?;
        Ok(result)
    }

    /// Execute one or more semicolon-separated SQL statements as a single
    /// atomic batch.  This is a convenience wrapper around
    /// [`execute_batch`](rusqlite::Connection::execute_batch) that wraps the
    /// statements in an explicit transaction so partial execution is never
    /// visible to other threads.
    ///
    /// # Example
    /// ```rust,no_run
    /// # use agentdb::AgentDB;
    /// let db = AgentDB::open(":memory:").unwrap();
    /// db.execute_batch(
    ///     "INSERT INTO _adb_nodes (id,kind,data,created_at,updated_at) VALUES ('a','t','{}',0,0);
    ///      INSERT INTO _adb_nodes (id,kind,data,created_at,updated_at) VALUES ('b','t','{}',0,0);"
    /// ).unwrap();
    /// ```
    pub fn execute_batch(&self, sql: &str) -> Result<()> {
        let mut conn = self.conn.lock().unwrap();
        let tx = conn.transaction()?;
        tx.execute_batch(sql)?;
        tx.commit()?;
        Ok(())
    }

    /// Query and return rows as JSON values
    pub fn query_json(&self, sql: &str) -> Result<Vec<serde_json::Value>> {
        let conn = self.conn.lock().unwrap();
        let mut stmt = conn.prepare(sql)?;
        let col_names: Vec<String> = stmt.column_names().iter().map(|s| s.to_string()).collect();
        let rows = stmt.query_map([], |row| {
            let mut map = serde_json::Map::new();
            for (i, name) in col_names.iter().enumerate() {
                let val: rusqlite::types::Value = row.get(i)?;
                map.insert(name.clone(), rusqlite_value_to_json(val));
            }
            Ok(serde_json::Value::Object(map))
        })?;
        rows.map(|r| r.map_err(crate::error::AgentDbError::Sqlite))
            .collect()
    }

    /// Query with parameters and return rows as JSON values.
    ///
    /// Use this for any query involving user-supplied values to prevent SQL injection.
    ///
    /// # Example
    /// ```rust,no_run
    /// # use agentdb::AgentDB;
    /// let db = AgentDB::open(":memory:").unwrap();
    /// let rows = db.query_json_params(
    ///     "SELECT * FROM _adb_nodes WHERE kind = ?1",
    ///     &[&"session" as &dyn rusqlite::ToSql],
    /// ).unwrap();
    /// ```
    pub fn query_json_params(
        &self,
        sql: &str,
        params: &[&dyn rusqlite::ToSql],
    ) -> Result<Vec<serde_json::Value>> {
        let conn = self.conn.lock().unwrap();
        let mut stmt = conn.prepare(sql)?;
        let col_names: Vec<String> = stmt.column_names().iter().map(|s| s.to_string()).collect();
        let rows = stmt.query_map(params, |row| {
            let mut map = serde_json::Map::new();
            for (i, name) in col_names.iter().enumerate() {
                let val: rusqlite::types::Value = row.get(i)?;
                map.insert(name.clone(), rusqlite_value_to_json(val));
            }
            Ok(serde_json::Value::Object(map))
        })?;
        rows.map(|r| r.map_err(crate::error::AgentDbError::Sqlite))
            .collect()
    }

    /// Flush dirty HNSW indexes and close gracefully.
    ///
    /// After this returns, the subsequent `Drop` is a no-op (no double flush).
    pub fn close(self) -> Result<()> {
        let collections = self.vectors().list_collections()?;
        for (name, dim, _) in collections {
            let col = self.vectors().collection(&name, dim)?;
            let is_dirty: i64 = {
                let conn = self.conn.lock().unwrap();
                conn.query_row(
                    "SELECT COALESCE(
                        (SELECT is_dirty FROM _adb_hnsw_index
                         WHERE collection_id =
                           (SELECT id FROM _adb_collections WHERE name = ?1)
                        ), 0)",
                    rusqlite::params![name],
                    |r| r.get(0),
                )
                .unwrap_or(0)
            };
            if is_dirty == 1 {
                col.reindex()?;
            }
        }
        // Checkpoint WAL to consolidate into a single file
        {
            let conn = self.conn.lock().unwrap();
            let _ = conn.pragma_update(None, "wal_checkpoint", "TRUNCATE");
        }
        self.closed.store(true, Ordering::Release);
        Ok(())
    }

    /// Return database-wide statistics (single-query implementation).
    pub fn stats(&self) -> Result<DbStats> {
        let conn = self.conn.lock().unwrap();
        conn.query_row(
            "SELECT
                 (SELECT COUNT(*)                FROM _adb_collections)     AS collections,
                 (SELECT COALESCE(SUM(count),0)  FROM _adb_collections)     AS vectors,
                 (SELECT COUNT(*)                FROM _adb_nodes)            AS nodes,
                 (SELECT COUNT(*)                FROM _adb_edges)            AS edges,
                 (SELECT COUNT(*)                FROM _adb_conversations)    AS conversations,
                 (SELECT COUNT(*)                FROM _adb_messages)         AS messages,
                 (SELECT COUNT(*)                FROM _adb_workflows)        AS workflows,
                 (SELECT COUNT(*)                FROM _adb_workflow_steps)   AS workflow_steps,
                 (SELECT COUNT(*)                FROM _adb_traces)           AS traces,
                 (SELECT COUNT(*)                FROM _adb_tools)            AS tools,
                 (SELECT COUNT(*)                FROM _adb_tool_calls)       AS tool_calls,
                 (SELECT COUNT(*)                FROM _adb_audit_log)        AS audit_entries,
                 (SELECT COUNT(*)                FROM _adb_prompt_templates) AS prompt_templates",
            [],
            |r| {
                Ok(DbStats {
                    collections: r.get(0)?,
                    vectors: r.get(1)?,
                    nodes: r.get(2)?,
                    edges: r.get(3)?,
                    conversations: r.get(4)?,
                    messages: r.get(5)?,
                    workflows: r.get(6)?,
                    workflow_steps: r.get(7)?,
                    traces: r.get(8)?,
                    tools: r.get(9)?,
                    tool_calls: r.get(10)?,
                    audit_entries: r.get(11)?,
                    prompt_templates: r.get(12)?,
                })
            },
        )
        .map_err(crate::error::AgentDbError::Sqlite)
    }
}

/// Database-wide statistics returned by [`AgentDB::stats`].
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct DbStats {
    /// Number of named vector collections.
    pub collections: i64,
    /// Total number of vectors across all collections.
    pub vectors: i64,
    /// Number of nodes in the memory graph.
    pub nodes: i64,
    /// Number of directed edges in the memory graph.
    pub edges: i64,
    /// Number of conversation threads.
    pub conversations: i64,
    /// Total number of messages across all conversations.
    pub messages: i64,
    /// Number of workflow records.
    pub workflows: i64,
    /// Total number of workflow steps across all workflows.
    pub workflow_steps: i64,
    /// Total number of reasoning trace entries.
    pub traces: i64,
    /// Number of registered tool definitions.
    pub tools: i64,
    /// Total number of logged tool calls.
    pub tool_calls: i64,
    /// Total number of audit log entries.
    pub audit_entries: i64,
    /// Total number of prompt template versions.
    pub prompt_templates: i64,
}

impl Drop for AgentDB {
    fn drop(&mut self) {
        if self.closed.load(Ordering::Acquire) {
            return;
        }
        if let Ok(collections) = self.vectors().list_collections() {
            for (name, dim, _) in collections {
                if let Ok(col) = self.vectors().collection(&name, dim) {
                    let is_dirty: bool = {
                        let conn = self.conn.lock().unwrap();
                        conn.query_row(
                            "SELECT COALESCE(
                                (SELECT is_dirty FROM _adb_hnsw_index
                                 WHERE collection_id =
                                   (SELECT id FROM _adb_collections WHERE name = ?1)
                                ), 0)",
                            rusqlite::params![name],
                            |r| r.get::<_, i64>(0),
                        )
                        .unwrap_or(0)
                            == 1
                    };
                    if is_dirty {
                        let _ = col.reindex();
                    }
                }
            }
        }
    }
}

// ── Internal connection accessor (used by sync module) ───────────────────────

impl AgentDB {
    /// Returns a clone of the internal `Arc<Mutex<Connection>>`.
    ///
    /// This is intentionally `pub(crate)` — only sub-modules of this crate
    /// that need to share the same SQLite connection (e.g. `sync`) should use
    /// it.  External callers should use the higher-level store accessors.
    pub(crate) fn conn_arc(&self) -> Arc<Mutex<Connection>> {
        Arc::clone(&self.conn)
    }
}

// ── WASM-only connection accessor ────────────────────────────────────────────

/// Methods only compiled when the `wasm` feature is active.
///
/// These provide low-level access to the underlying `rusqlite::Connection`
/// required by the OPFS persistence layer (`crate::wasm_opfs`). They are
/// intentionally not exposed in non-WASM builds to avoid leaking internals.
#[cfg(feature = "wasm")]
impl AgentDB {
    /// Run a closure with exclusive access to the underlying SQLite connection.
    ///
    /// Used by `wasm_opfs::sqlite_serialize` and `sqlite_deserialize` which
    /// need the raw `*mut sqlite3` handle.
    ///
    /// The closure **must not** panic while holding the lock.
    pub(crate) fn with_conn<F, T>(&self, f: F) -> crate::error::Result<T>
    where
        F: FnOnce(&rusqlite::Connection) -> crate::error::Result<T>,
    {
        let conn = self.conn.lock().unwrap();
        f(&conn)
    }
}

fn rusqlite_value_to_json(val: rusqlite::types::Value) -> serde_json::Value {
    match val {
        rusqlite::types::Value::Null => serde_json::Value::Null,
        rusqlite::types::Value::Integer(i) => serde_json::Value::Number(i.into()),
        rusqlite::types::Value::Real(f) => serde_json::Number::from_f64(f)
            .map(serde_json::Value::Number)
            .unwrap_or(serde_json::Value::Null),
        rusqlite::types::Value::Text(s) => serde_json::Value::String(s),
        rusqlite::types::Value::Blob(b) => {
            serde_json::Value::String(format!("<blob {} bytes>", b.len()))
        }
    }
}