llm_kernel/graph/

backend.rs

//! Backend-agnostic graph trait and SQLite implementation.
//!
//! [`GraphBackend`] is a sync, object-safe trait covering the primitive node/edge
//! operations every graph backend must support. It deliberately exposes **no
//! `rusqlite` types**, so a PostgreSQL or in-memory backend can implement it
//! (see the v0.8.0 roadmap). [`SqliteGraph`] is the bundled implementation: it
//! wraps a single mutex-guarded connection and delegates to the existing
//! free-function graph API in [`crate::graph`].
//!
//! The existing free functions (`upsert_node(&conn, …)`, `search_nodes(&conn, …)`,
//! …) are unchanged — callers that already own a `Connection` keep using them.
//! [`SqliteGraph`] simply packages a connection behind the trait for users who
//! want backend-agnostic graph access.

use std::path::Path;
use std::sync::Mutex;

use rusqlite::Connection;

use crate::error::Result;
use crate::graph::recall::smart_recall;
use crate::graph::schema::{init_graph_schema, migrate_graph, schema_version};
use crate::graph::search::{query_nodes, search_nodes};
use crate::graph::store::{
    append_edge, delete_edge, delete_node, edges_for_node, remove_edges_for_node, upsert_node,
};
use crate::graph::traversal::related_nodes;
use crate::graph::types::{EdgeDirection, GraphEdge, GraphNode, ScoredNode};

/// Sync, object-safe trait for graph backends.
///
/// Methods cover node/edge CRUD, FTS search, filtered query, and schema
/// migration. No method exposes `rusqlite` types, so the trait is implementable
/// by any backend. `dyn GraphBackend` is usable.
pub trait GraphBackend: Send + Sync {
    /// Insert or replace a node.
    fn upsert_node(&self, node: &GraphNode) -> Result<()>;
    /// Read a single node by ID (`None` if absent).
    fn read_node(&self, id: &str) -> Result<Option<GraphNode>>;
    /// Delete a node by ID. Returns `true` if a row was removed.
    fn delete_node(&self, id: &str) -> Result<bool>;
    /// FTS5 full-text search, ranked by importance DESC.
    fn search_nodes(&self, query: &str, limit: usize) -> Result<Vec<GraphNode>>;
    /// Dynamic filter by tag / node_type / project.
    #[allow(clippy::too_many_arguments)]
    fn query_nodes(
        &self,
        tag: Option<&str>,
        node_type: Option<&str>,
        project: Option<&str>,
        limit: usize,
    ) -> Result<Vec<GraphNode>>;
    /// Composite recall — rank nodes by recency, importance, access, FTS, and
    /// graph boost. The canonical high-level read path for "what's relevant".
    fn smart_recall(
        &self,
        project: Option<&str>,
        hint: Option<&str>,
        limit: usize,
    ) -> Result<Vec<ScoredNode>>;
    /// BFS-traverse up to `depth` hops from `start_id`, returning related node
    /// IDs (excluding the start).
    fn related_nodes(&self, start_id: &str, depth: usize) -> Result<Vec<String>>;
    /// Append an edge (duplicates by edge ID are ignored).
    fn append_edge(&self, edge: &GraphEdge) -> Result<()>;
    /// Append many edges in one call.
    ///
    /// Duplicates by edge ID *or* by the `(source, target, relation)` unique
    /// index are ignored. The default implementation loops [`Self::append_edge`];
    /// backends with a batch path override it for throughput (citation graphs
    /// built during indexing can reach hundreds of thousands of edges).
    fn append_edges(&self, edges: &[GraphEdge]) -> Result<()> {
        for edge in edges {
            self.append_edge(edge)?;
        }
        Ok(())
    }
    /// Read edges touching `node_id`, filtered by direction and an optional
    /// relation. The default implementation reads both directions via
    /// [`Self::edges_for_node`] and filters in Rust; backends with directional
    /// indexes override for efficiency.
    fn edges_for_node_dir(
        &self,
        node_id: &str,
        dir: EdgeDirection,
        relation: Option<&str>,
    ) -> Result<Vec<GraphEdge>> {
        Ok(self
            .edges_for_node(node_id)?
            .into_iter()
            .filter(|e| match dir {
                EdgeDirection::Out => e.source == node_id,
                EdgeDirection::In => e.target == node_id,
                EdgeDirection::Both => true,
            })
            .filter(|e| relation.is_none_or(|r| e.relation == r))
            .collect())
    }
    /// 1-hop neighbors of `seed_ids` (weighted sum), restricted by direction
    /// and an optional relation. Seed nodes are excluded. The default
    /// implementation walks each seed via [`Self::edges_for_node_dir`].
    fn neighbors_weighted(
        &self,
        seed_ids: &[String],
        dir: EdgeDirection,
        relation: Option<&str>,
    ) -> Result<Vec<(String, f64)>> {
        let mut weights: std::collections::HashMap<String, f64> = std::collections::HashMap::new();
        let seed_set: std::collections::HashSet<&str> =
            seed_ids.iter().map(String::as_str).collect();
        for seed in seed_ids {
            for e in self.edges_for_node_dir(seed, dir, relation)? {
                let other = if e.source == *seed {
                    &e.target
                } else {
                    &e.source
                };
                if !seed_set.contains(other.as_str()) {
                    *weights.entry(other.clone()).or_default() += e.weight;
                }
            }
        }
        let mut result: Vec<(String, f64)> = weights.into_iter().collect();
        result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        Ok(result)
    }
    /// BFS-traverse up to `depth` hops from `start_id`, returning related node
    /// IDs (excluding the start), restricted by direction and an optional
    /// relation. The default implementation hops via [`Self::neighbors_weighted`].
    fn related_nodes_filtered(
        &self,
        start_id: &str,
        depth: usize,
        dir: EdgeDirection,
        relation: Option<&str>,
    ) -> Result<Vec<String>> {
        if depth == 0 {
            return Ok(vec![]);
        }
        let mut visited: std::collections::HashSet<String> = std::collections::HashSet::new();
        visited.insert(start_id.to_string());
        let mut frontier: Vec<String> = vec![start_id.to_string()];
        for _ in 0..depth {
            let mut next: Vec<String> = Vec::new();
            for node in &frontier {
                for (nb, _) in self.neighbors_weighted(std::slice::from_ref(node), dir, relation)? {
                    if visited.insert(nb.clone()) {
                        next.push(nb);
                    }
                }
            }
            if next.is_empty() {
                break;
            }
            frontier = next;
        }
        visited.remove(start_id);
        Ok(visited.into_iter().collect())
    }
    /// Read edges where the given node is source or target.
    fn edges_for_node(&self, node_id: &str) -> Result<Vec<GraphEdge>>;
    /// Delete an edge by ID. Returns `true` if a row was removed.
    fn delete_edge(&self, id: &str) -> Result<bool>;
    /// Remove every edge connected to a node.
    fn remove_edges_for_node(&self, node_id: &str) -> Result<()>;

    /// Recorded schema version for this backend.
    fn current_version(&self) -> Result<u32>;
    /// Apply pending migrations up to the backend's latest schema version.
    /// Returns the resulting version.
    fn migrate(&self) -> Result<u32>;
}

/// SQLite-backed [`GraphBackend`] over one mutex-guarded connection.
///
/// Opening applies the schema and runs any pending migrations, so a database
/// created by an older `llm-kernel` is upgraded transparently on open.
pub struct SqliteGraph {
    conn: Mutex<Connection>,
}

impl SqliteGraph {
    /// Open (or create) a graph database at `path`, applying schema + migrations.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let conn = open_with_schema(path.as_ref())?;
        Ok(Self {
            conn: Mutex::new(conn),
        })
    }

    /// Create an in-memory graph (useful for tests and ephemeral stores).
    pub fn open_in_memory() -> Result<Self> {
        let conn = Connection::open_in_memory().map_err(store_err)?;
        init_graph_schema(&conn)?;
        let current = schema_version(&conn)?;
        migrate_graph(&conn, current)?;
        Ok(Self {
            conn: Mutex::new(conn),
        })
    }

    /// Lock helper: recover the guard even if a previous holder panicked.
    fn lock(&self) -> std::sync::MutexGuard<'_, Connection> {
        self.conn.lock().unwrap_or_else(|e| e.into_inner())
    }

    /// Record that a node's content was verified as of `now` (ISO 8601).
    ///
    /// Convenience wrapper over [`crate::graph::lifecycle::mark_verified`] —
    /// the connection is owned by the graph, so callers don't need to open
    /// their own.
    pub fn mark_verified(&self, id: &str, now: &str) -> Result<bool> {
        let c = self.lock();
        crate::graph::lifecycle::mark_verified(&c, id, now)
    }

    /// Count nodes whose `valid_until` is set and earlier than `now` (ISO 8601).
    ///
    /// Convenience wrapper over [`crate::graph::lifecycle::count_expired_nodes`].
    pub fn count_expired_nodes(&self, now: &str) -> Result<u64> {
        let c = self.lock();
        crate::graph::lifecycle::count_expired_nodes(&c, now)
    }

    /// Run `f` against the underlying connection inside one transaction.
    ///
    /// Commits when `f` returns `Ok`, rolls back on `Err` — so multi-step
    /// sequences (e.g. delete-then-insert edge replacement) cannot leave a
    /// half-applied state. `f` receives `&Connection`: the transaction derefs
    /// to it, so the free functions in [`crate::graph::store`] work as-is.
    ///
    /// Do **not** call other `SqliteGraph` methods on the same instance inside
    /// `f` — they re-acquire the connection lock and will deadlock. Use only
    /// the free functions on the passed connection.
    pub fn with_tx(&self, f: impl FnOnce(&Connection) -> Result<()>) -> Result<()> {
        let conn = self.lock();
        let tx = conn.unchecked_transaction().map_err(store_err)?;
        f(&tx)?;
        tx.commit().map_err(store_err)
    }
}

/// Open a file-backed connection, apply the schema, then run pending migrations.
fn open_with_schema(path: &Path) -> Result<Connection> {
    let conn = Connection::open(path).map_err(store_err)?;
    init_graph_schema(&conn)?;
    let current = schema_version(&conn)?;
    migrate_graph(&conn, current)?;
    Ok(conn)
}

fn store_err(e: rusqlite::Error) -> crate::error::KernelError {
    crate::error::KernelError::Store(e.to_string())
}

impl GraphBackend for SqliteGraph {
    fn upsert_node(&self, node: &GraphNode) -> Result<()> {
        let c = self.lock();
        upsert_node(&c, node)
    }

    fn read_node(&self, id: &str) -> Result<Option<GraphNode>> {
        let c = self.lock();
        crate::graph::store::read_node(&c, id)
    }

    fn delete_node(&self, id: &str) -> Result<bool> {
        let c = self.lock();
        delete_node(&c, id)
    }

    fn search_nodes(&self, query: &str, limit: usize) -> Result<Vec<GraphNode>> {
        let c = self.lock();
        search_nodes(&c, query, limit)
    }

    fn query_nodes(
        &self,
        tag: Option<&str>,
        node_type: Option<&str>,
        project: Option<&str>,
        limit: usize,
    ) -> Result<Vec<GraphNode>> {
        let c = self.lock();
        query_nodes(&c, tag, node_type, project, limit)
    }

    fn smart_recall(
        &self,
        project: Option<&str>,
        hint: Option<&str>,
        limit: usize,
    ) -> Result<Vec<ScoredNode>> {
        let c = self.lock();
        smart_recall(&c, project, hint, limit)
    }

    fn related_nodes(&self, start_id: &str, depth: usize) -> Result<Vec<String>> {
        let c = self.lock();
        Ok(related_nodes(&c, start_id, depth))
    }

    fn append_edge(&self, edge: &GraphEdge) -> Result<()> {
        let c = self.lock();
        append_edge(&c, edge)
    }

    fn append_edges(&self, edges: &[GraphEdge]) -> Result<()> {
        let c = self.lock();
        crate::graph::store::append_edges(&c, edges)
    }

    fn edges_for_node_dir(
        &self,
        node_id: &str,
        dir: EdgeDirection,
        relation: Option<&str>,
    ) -> Result<Vec<GraphEdge>> {
        let c = self.lock();
        crate::graph::store::edges_for_node_dir(&c, node_id, dir, relation)
    }

    fn neighbors_weighted(
        &self,
        seed_ids: &[String],
        dir: EdgeDirection,
        relation: Option<&str>,
    ) -> Result<Vec<(String, f64)>> {
        let c = self.lock();
        Ok(crate::graph::traversal::neighbors_weighted(
            &c, seed_ids, dir, relation,
        ))
    }

    fn edges_for_node(&self, node_id: &str) -> Result<Vec<GraphEdge>> {
        let c = self.lock();
        edges_for_node(&c, node_id)
    }

    fn delete_edge(&self, id: &str) -> Result<bool> {
        let c = self.lock();
        delete_edge(&c, id)
    }

    fn remove_edges_for_node(&self, node_id: &str) -> Result<()> {
        let c = self.lock();
        remove_edges_for_node(&c, node_id)
    }

    fn current_version(&self) -> Result<u32> {
        let c = self.lock();
        schema_version(&c)
    }

    fn migrate(&self) -> Result<u32> {
        let c = self.lock();
        let current = schema_version(&c)?;
        migrate_graph(&c, current)
    }
}

/// CJK-aware convenience methods (available with the `graph-cjk` feature).
#[cfg(feature = "graph-cjk")]
impl SqliteGraph {
    /// CJK (or mixed) search via contiguous substring matching.
    ///
    /// Delegates to [`crate::graph::cjk::search_nodes_cjk`]; see its docs for
    /// the matching semantics.
    pub fn search_nodes_cjk(&self, query: &str, limit: usize) -> Result<Vec<GraphNode>> {
        let c = self.lock();
        crate::graph::cjk::search_nodes_cjk(&c, query, limit)
    }

    /// Segment a string for CJK tokenization (exposed for callers that want to
    /// pre-process queries or inspect tokenization).
    pub fn segment_cjk(text: &str) -> String {
        crate::graph::cjk::segment_cjk(text)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn sample_node(id: &str) -> GraphNode {
        GraphNode {
            id: id.to_string(),
            node_type: "concept".to_string(),
            title: format!("Node {id}"),
            body: "graph backend test body".to_string(),
            tags: vec!["backend".to_string()],
            projects: vec![],
            agents: vec![],
            created: "2026-01-01T00:00:00Z".to_string(),
            updated: "2026-01-01T00:00:00Z".to_string(),
            importance: 0.5,
            access_count: 0,
            accessed_at: String::new(),
            ..Default::default()
        }
    }

    /// AC4: a node round-trips through the trait, and the trait is usable as
    /// `dyn GraphBackend` (object-safety) with no `rusqlite` in the surface.
    #[test]
    fn dyn_backend_round_trips_node() {
        let backend: Box<dyn GraphBackend> = Box::new(SqliteGraph::open_in_memory().unwrap());
        assert!(backend.read_node("n1").unwrap().is_none());
        backend.upsert_node(&sample_node("n1")).unwrap();
        let loaded = backend.read_node("n1").unwrap().unwrap();
        assert_eq!(loaded.title, "Node n1");
        assert_eq!(loaded.tags, vec!["backend".to_string()]);
        assert!(backend.delete_node("n1").unwrap());
        assert!(backend.read_node("n1").unwrap().is_none());
    }

    /// Temporal-validity wrappers reachable without owning a connection.
    #[test]
    fn mark_verified_and_count_expired_wrappers() {
        let backend = SqliteGraph::open_in_memory().unwrap();
        backend.upsert_node(&sample_node("n1")).unwrap();
        assert!(backend.mark_verified("n1", "2026-08-18T00:00:00Z").unwrap());
        let node = backend.read_node("n1").unwrap().unwrap();
        assert_eq!(node.last_verified, "2026-08-18T00:00:00Z");
        assert_eq!(
            backend.count_expired_nodes("2026-08-18T00:00:00Z").unwrap(),
            0
        );
    }

    /// `with_tx` commits on Ok — both statements are visible afterwards.
    #[test]
    fn with_tx_commits_on_ok() {
        let backend = SqliteGraph::open_in_memory().unwrap();
        backend.upsert_node(&sample_node("a")).unwrap();
        backend.upsert_node(&sample_node("b")).unwrap();
        backend
            .with_tx(|conn| {
                crate::graph::store::remove_edges_for_node(conn, "a")?;
                crate::graph::store::append_edge(
                    conn,
                    &GraphEdge {
                        id: "e1".into(),
                        source: "a".into(),
                        target: "b".into(),
                        relation: "related".into(),
                        weight: 1.0,
                        ts: "2026-01-01T00:00:00Z".into(),
                    },
                )
            })
            .unwrap();
        let edges = backend.edges_for_node("a").unwrap();
        assert_eq!(edges.len(), 1);
    }

    /// `with_tx` rolls back on Err — the delete is undone, nothing is lost.
    #[test]
    fn with_tx_rolls_back_on_err() {
        let backend = SqliteGraph::open_in_memory().unwrap();
        backend.upsert_node(&sample_node("a")).unwrap();
        backend.upsert_node(&sample_node("b")).unwrap();
        backend
            .append_edge(&GraphEdge {
                id: "e0".into(),
                source: "a".into(),
                target: "b".into(),
                relation: "related".into(),
                weight: 1.0,
                ts: "2026-01-01T00:00:00Z".into(),
            })
            .unwrap();

        let result = backend.with_tx(|conn| {
            crate::graph::store::remove_edges_for_node(conn, "a")?;
            // Simulate a failure after the delete.
            Err(crate::error::KernelError::Store("boom".into()))
        });
        assert!(result.is_err());
        // The original edge survived — rollback undid the delete.
        let edges = backend.edges_for_node("a").unwrap();
        assert_eq!(edges.len(), 1, "edge lost mid-transaction: {edges:?}");
    }

    /// AC5: a fresh backend reports the current schema version.
    #[test]
    fn fresh_backend_reports_current_version() {
        let backend = SqliteGraph::open_in_memory().unwrap();
        assert_eq!(
            backend.current_version().unwrap(),
            crate::graph::schema::GRAPH_SCHEMA_VERSION
        );
    }

    /// AC5: search through the trait finds an inserted node by title.
    #[test]
    fn backend_search_finds_node() {
        let backend = SqliteGraph::open_in_memory().unwrap();
        backend.upsert_node(&sample_node("rust")).unwrap();
        let hits = backend.search_nodes("graph backend", 10).unwrap();
        assert_eq!(hits.len(), 1);
        assert_eq!(hits[0].id, "rust");
    }

    /// The composite recall path is reachable through the trait.
    #[test]
    fn backend_smart_recall_finds_relevant() {
        let backend = SqliteGraph::open_in_memory().unwrap();
        let mut n = sample_node("rust");
        n.body = "rust ownership borrow checker".to_string();
        backend.upsert_node(&n).unwrap();
        let recalled = backend.smart_recall(None, Some("ownership"), 5).unwrap();
        assert!(recalled.iter().any(|s| s.node.id == "rust"));
    }

    /// The composite traversal path is reachable through the trait.
    #[test]
    fn backend_related_nodes_traverses_edges() {
        let backend = SqliteGraph::open_in_memory().unwrap();
        backend.upsert_node(&sample_node("a")).unwrap();
        backend.upsert_node(&sample_node("b")).unwrap();
        backend
            .append_edge(&GraphEdge {
                id: "e1".into(),
                source: "a".into(),
                target: "b".into(),
                relation: "related".into(),
                weight: 1.0,
                ts: "2026-01-01T00:00:00Z".into(),
            })
            .unwrap();
        let related = backend.related_nodes("a", 2).unwrap();
        assert!(related.contains(&"b".to_string()));
    }

    /// Batch edge append + directional/relation-filtered lookups through the
    /// trait (object-safe path), including the BFS `related_nodes_filtered`.
    #[test]
    fn dyn_backend_batch_and_filtered_edges() {
        let backend: Box<dyn GraphBackend> = Box::new(SqliteGraph::open_in_memory().unwrap());
        backend
            .append_edges(&[
                GraphEdge {
                    id: "e1".into(),
                    source: "a".into(),
                    target: "b".into(),
                    relation: "cites".into(),
                    weight: 1.0,
                    ts: "t".into(),
                },
                GraphEdge {
                    id: "e2".into(),
                    source: "c".into(),
                    target: "a".into(),
                    relation: "cites".into(),
                    weight: 1.0,
                    ts: "t".into(),
                },
                GraphEdge {
                    id: "e3".into(),
                    source: "a".into(),
                    target: "d".into(),
                    relation: "see_also".into(),
                    weight: 1.0,
                    ts: "t".into(),
                },
            ])
            .unwrap();
        // Out-edges of `a`: b, d.
        assert_eq!(
            backend
                .edges_for_node_dir("a", EdgeDirection::Out, None)
                .unwrap()
                .len(),
            2
        );
        // Out-only neighbors of `a` filtered to `cites`: only b (not c, which is in-edge).
        let nbs = backend
            .neighbors_weighted(&["a".to_string()], EdgeDirection::Out, Some("cites"))
            .unwrap();
        let ids: Vec<&str> = nbs.iter().map(|(id, _)| id.as_str()).collect();
        assert_eq!(ids, vec!["b"]);
        // BFS out-only depth 2 from `a`: reaches b and d.
        let rel = backend
            .related_nodes_filtered("a", 2, EdgeDirection::Out, None)
            .unwrap();
        assert!(rel.contains(&"b".to_string()));
        assert!(rel.contains(&"d".to_string()));
    }
}