gonzalo_knowledge/

lib.rs

//! Knowledge-store capability for gonzalo.
//!
//! A single "what do we know about X" surface that composes the existing
//! capability layers (ADR 0011): a [`Store`] for records, a [`VectorIndex`] for
//! semantic retrieval, and an [`Embedder`] for turning text into vectors — all
//! addressed by the shared [`RecordKey`](gonzalo_core::RecordKey). Queries return first-class
//! [`Record`]s as [`Hit`]s, not bare ids.
//!
//! Which record kinds are knowledge-bearing, and how their text is split into
//! embeddable pieces, lives in [`chunk`] (with [`knowledge_text`] as the joined
//! one-document view). Records are embedded at chunk granularity — a `Session`
//! by turn, a long `MemoryTier`/`Ticket` by paragraph, a `Topic` by bullet — so
//! a query matching one turn/section isn't drowned out by a document average;
//! [`KnowledgeStore::query`] de-dups chunk matches back to one hit per record.
//! A `graph`-backed join ([`KnowledgeStore::query_in_subgraph`]) is a further
//! refinement behind this same surface.

use gonzalo_core::{KeyPrefix, Record, RecordKind, Result, Store};
use gonzalo_domain::{MemoryTier, RecordCodec, Session, Ticket, TicketEvent, Topic};
use gonzalo_vector::{Embedder, VectorIndex};

/// One search hit: a first-class record and its similarity score (higher is
/// more similar).
#[derive(Debug, Clone)]
pub struct Hit {
    pub record: Record,
    pub score: f32,
}

/// One chunk-granular search hit: the parent [`Record`], the matching chunk's
/// score, its `ordinal` within the record, and that chunk's `text`. Unlike a
/// [`Hit`], chunk hits are **not** de-duped to one-per-record.
#[derive(Debug, Clone)]
pub struct ChunkHit {
    pub record: Record,
    pub score: f32,
    pub ordinal: usize,
    pub text: String,
}

/// Composes a [`Store`], a [`VectorIndex`], and an [`Embedder`] into one
/// retrieval surface keyed by [`RecordKey`](gonzalo_core::RecordKey).
pub struct KnowledgeStore<S, V, E> {
    store: S,
    index: V,
    embedder: E,
    /// Last-seen chunk count per record, so a re-ingest that shrinks a record
    /// can remove the now-orphaned high-ordinal chunks from the index. Always
    /// held in memory, so [`new`](Self::new) starts it empty — correct for a
    /// fresh index with nothing indexed yet. [`open`](Self::open) rebuilds it
    /// instead by scanning a durable index's existing keys, so counts survive
    /// a restart even though this field itself does not.
    chunk_counts: std::sync::Mutex<std::collections::HashMap<gonzalo_core::RecordKey, usize>>,
}

impl<S: Store, V: VectorIndex, E: Embedder> KnowledgeStore<S, V, E> {
    pub fn new(store: S, index: V, embedder: E) -> Self {
        Self {
            store,
            index,
            embedder,
            chunk_counts: std::sync::Mutex::new(std::collections::HashMap::new()),
        }
    }

    /// Open over an index that may already hold vectors, rebuilding the chunk
    /// counts from it.
    ///
    /// [`new`](Self::new) starts those counts empty, which is right for a fresh
    /// in-memory index and wrong for a durable one: a re-ingest that shrinks a
    /// record would not know how many chunks to remove, leaving orphans that
    /// still match queries (#150). `KeyPrefix` cannot narrow below a collection,
    /// so this scans once here rather than per ingest.
    pub async fn open(store: S, index: V, embedder: E) -> Result<Self> {
        let mut counts: std::collections::HashMap<gonzalo_core::RecordKey, usize> =
            std::collections::HashMap::new();
        for chunk in index.keys(&KeyPrefix::default()).await? {
            let (parent, ordinal) = parent_key(&chunk);
            let entry = counts.entry(parent).or_insert(0);
            *entry = (*entry).max(ordinal + 1);
        }
        Ok(Self {
            store,
            index,
            embedder,
            chunk_counts: std::sync::Mutex::new(counts),
        })
    }

    /// Borrow the underlying store (e.g. to put records before ingesting them).
    pub fn store(&self) -> &S {
        &self.store
    }

    /// Ingest the record at `key`: split it into per-kind [`chunk`]s, embed each,
    /// and index it under a derived per-chunk key. Returns `Ok(false)` (without
    /// indexing) if the record is absent or its kind is definitionally not
    /// knowledge-bearing; returns `Err` if a knowledge-bearing body fails to
    /// parse — a corrupt record is a real failure, not "not indexable" (#139).
    ///
    /// **One batch per record, not one call per chunk** (gonzalo#326). On a
    /// durable index every `upsert` is a full OCC commit — re-serialise the
    /// shard, upload it as a blob, write a new manifest revision — so a
    /// three-chunk record used to advance the manifest three times and leave two
    /// superseded shard blobs for `gc`. Every chunk is now embedded first, then
    /// written with one [`VectorIndex::upsert_many`]. A consequence of embedding
    /// first: if an embedding fails part way, nothing is written at all, where
    /// before the earlier chunks were already committed.
    ///
    /// The orphan removals below are **not** batched yet. Both candidate shapes
    /// — a `remove_many` on the trait, or folding removals into the upsert
    /// commit — add a `VectorIndex` method, and `cargo publish --workspace
    /// --dry-run` verifies each packaged crate against the *published*
    /// `gonzalo-vector`, so a new trait method cannot be called from this crate
    /// until the workspace version is bumped. Orphans only occur when a record
    /// shrinks, so the common path commits once either way. Tracked separately.
    pub async fn ingest(&self, key: &gonzalo_core::RecordKey) -> Result<bool> {
        let Some(record) = self.store.get(key).await? else {
            return Ok(false);
        };
        let Some(chunks) = chunk(&record)? else {
            return Ok(false);
        };

        // Remove chunks orphaned by a shrink since the last ingest. Read the old
        // count under a scoped lock — never held across an `.await`.
        let old = {
            let counts = self.chunk_counts.lock().unwrap();
            counts.get(key).copied().unwrap_or(0)
        };
        for ordinal in chunks.len()..old {
            self.index.remove(&chunk_key(key, ordinal)).await?;
        }

        let mut items = Vec::with_capacity(chunks.len());
        for (ordinal, text) in chunks.iter().enumerate() {
            items.push((chunk_key(key, ordinal), self.embedder.embed(text).await?));
        }
        if !items.is_empty() {
            self.index.upsert_many(items).await?;
        }

        self.chunk_counts
            .lock()
            .unwrap()
            .insert(key.clone(), chunks.len());
        Ok(true)
    }

    /// De-index the record at `key`: remove every chunk vector it contributed to
    /// the index and forget its chunk count. This is the counterpart to a record
    /// being deleted from the store — without it, the record's chunk vectors are
    /// orphaned, still matching in [`query`] but de-duping to a parent that no
    /// longer resolves, so they silently shrink the result set (#150).
    ///
    /// Idempotent: removing a record that was never ingested (chunk count 0)
    /// removes nothing. Uses the same [`chunk_key`] derivation as [`ingest`].
    pub async fn remove(&self, key: &gonzalo_core::RecordKey) -> Result<()> {
        // Read the count under a scoped lock — never held across an `.await`.
        let count = {
            let counts = self.chunk_counts.lock().unwrap();
            counts.get(key).copied().unwrap_or(0)
        };
        for ordinal in 0..count {
            self.index.remove(&chunk_key(key, ordinal)).await?;
        }
        self.chunk_counts.lock().unwrap().remove(key);
        Ok(())
    }

    /// Semantic query, one hit per record. Embed `text`, over-fetch chunk
    /// matches (restricted to `filter`), collapse them to their parent records
    /// keeping each record's best chunk score, and return the top-`k` records.
    ///
    /// Because a record with many matching chunks can crowd the over-fetch
    /// window, pathological cases may return fewer than `k` hits.
    pub async fn query(&self, text: &str, k: usize, filter: &KeyPrefix) -> Result<Vec<Hit>> {
        if k == 0 {
            return Ok(Vec::new());
        }
        let query_vec = self.embedder.embed(text).await?;
        let fetch = k.saturating_mul(OVERFETCH);
        let matches = self.index.query(&query_vec, fetch, filter).await?;

        // Collapse chunk matches to parents, keeping the best score per parent.
        let mut best: std::collections::BTreeMap<gonzalo_core::RecordKey, f32> =
            std::collections::BTreeMap::new();
        for m in matches {
            let (parent, _ordinal) = parent_key(&m.key);
            best.entry(parent)
                .and_modify(|s| *s = s.max(m.score))
                .or_insert(m.score);
        }

        // Rank parents by best chunk score, resolve the top-k to records.
        let mut ranked: Vec<(gonzalo_core::RecordKey, f32)> = best.into_iter().collect();
        ranked.sort_by(|a, b| b.1.total_cmp(&a.1));
        ranked.truncate(k);

        let mut hits = Vec::with_capacity(ranked.len());
        for (parent, score) in ranked {
            if let Some(record) = self.store.get(&parent).await? {
                hits.push(Hit { record, score });
            }
        }
        Ok(hits)
    }

    /// Chunk-granular query: the top-`k` matching chunks (restricted to
    /// `filter`), **not** de-duped to their parents. Each hit carries the parent
    /// record, the chunk's ordinal, and the chunk's text.
    pub async fn query_chunks(
        &self,
        text: &str,
        k: usize,
        filter: &KeyPrefix,
    ) -> Result<Vec<ChunkHit>> {
        if k == 0 {
            return Ok(Vec::new());
        }
        let query_vec = self.embedder.embed(text).await?;
        let matches = self.index.query(&query_vec, k, filter).await?;
        let mut hits = Vec::with_capacity(matches.len());
        for m in matches {
            let (parent, ordinal) = parent_key(&m.key);
            if let Some(record) = self.store.get(&parent).await? {
                let text = chunk(&record)?
                    .and_then(|cs| cs.into_iter().nth(ordinal))
                    .unwrap_or_default();
                hits.push(ChunkHit {
                    record,
                    score: m.score,
                    ordinal,
                    text,
                });
            }
        }
        Ok(hits)
    }

    /// **Vector⋈graph** (ADR 0011): rank the top-`k` candidates by semantic
    /// similarity, then keep only those whose record falls in `root`'s
    /// call-graph neighborhood in `graph` — `root` itself, its callers, and its
    /// callees. Composition is by shared key: a record is in the neighborhood
    /// when its [`RecordKey`](gonzalo_core::RecordKey)`.id` equals a symbol name
    /// in the neighborhood (no new cross-store linkage, ADR 0008/0011).
    ///
    /// So "semantically similar to X **and** structurally near `root`" is one
    /// call. The result is the in-neighborhood subset of the top-`k` (≤ `k`).
    #[cfg(feature = "graph")]
    pub async fn query_in_subgraph(
        &self,
        text: &str,
        k: usize,
        graph: &dyn gonzalo_graph::GraphStore,
        root: &str,
    ) -> Result<Vec<Hit>> {
        use std::collections::BTreeSet;
        let mut neighborhood: BTreeSet<String> = BTreeSet::from([root.to_string()]);
        neighborhood.extend(graph.callers_of(root));
        neighborhood.extend(graph.callees(root));

        let hits = self.query(text, k, &KeyPrefix::default()).await?;
        Ok(hits
            .into_iter()
            .filter(|h| neighborhood.contains(&h.record.key.id))
            .collect())
    }
}

/// The embeddable text for a record, or `None` if its kind is not
/// knowledge-bearing (ADR 0011). Kept as the one-document view: the [`chunk`]s
/// joined back into a single string.
pub fn knowledge_text(record: &Record) -> Option<String> {
    // A body that fails to parse is not a "one document"; collapse it to `None`
    // for this convenience view. The indexing path ([`chunk`]/[`ingest`])
    // surfaces the error instead (#139).
    chunk(record).ok().flatten().map(|cs| cs.join("\n"))
}

/// Split a record into ordered chunk texts. Returns `Ok(None)` if the kind is
/// definitionally not knowledge-bearing (`Checkpoint`/`GraphManifest`, ADR
/// 0011), and `Err` if a knowledge-bearing body fails to parse — so a corrupt
/// record is distinguishable from a genuinely non-indexable one (#139).
/// Chunking is per-kind: a `Session` by turn, a long `MemoryTier`/`Ticket` by
/// section/paragraph, a `Topic` by bullet. The kind's title/name is folded into
/// the first chunk so it stays searchable. A kind that yields a single piece is
/// the one-document fallback — identical to phase-1 behavior. Extraction goes
/// through the `gonzalo-domain` typed views.
pub fn chunk(record: &Record) -> Result<Option<Vec<String>>> {
    let chunks = match record.kind {
        RecordKind::MemoryTier => {
            let t = MemoryTier::from_body(&record.body)?;
            prepend_header(&t.name, split_paragraphs(&t.content))
        }
        RecordKind::Topic => {
            let t = Topic::from_body(&record.body)?;
            prepend_header(&t.slug, t.bullets)
        }
        RecordKind::Session => {
            let s = Session::from_body(&record.body)?;
            prepend_header(&s.name, s.turns.into_iter().map(|turn| turn.text).collect())
        }
        RecordKind::Ticket => {
            let t = Ticket::from_body(&record.body)?;
            let mut cs = vec![format!("{}\n{}", t.title, t.labels.join(" "))];
            cs.extend(split_paragraphs(&t.body.markdown));
            cs
        }
        RecordKind::TicketEvent => {
            let e = TicketEvent::from_body(&record.body)?;
            vec![e.body]
        }
        // Not knowledge-bearing: a checkpoint is opaque state; a graph manifest
        // is a path -> content-hash map, not natural-language text (ADR 0011/0012).
        // Checkpoint and GraphManifest are the ONLY definitionally-opaque
        // kinds — a parse failure on a knowledge-bearing kind above
        // propagates as an error rather than being silently indistinguishable
        // from "not indexable" (#139).
        // Tombstone is listed here defensively, not definitionally: consumer
        // reads never surface one, so indexing never actually sees this kind
        // (ADR 0021).
        // Fleet access-control records are configuration and audit data, and
        // hold personal data (emails, handles), so they are never indexed for
        // semantic search (ADR 0022).
        RecordKind::Checkpoint
        | RecordKind::GraphManifest
        | RecordKind::VectorManifest
        | RecordKind::Tombstone
        | RecordKind::Person
        | RecordKind::IdentityBinding
        | RecordKind::RoleGrant
        | RecordKind::ChannelConfig
        | RecordKind::LinkToken
        | RecordKind::AuditEntry => {
            return Ok(None);
        }
    };
    Ok(Some(chunks))
}

/// Fold `header` into the first `pieces` chunk so it stays searchable; if
/// `pieces` is empty, the header stands alone as the single chunk.
fn prepend_header(header: &str, mut pieces: Vec<String>) -> Vec<String> {
    match pieces.first_mut() {
        Some(first) => *first = format!("{header}\n{first}"),
        None => pieces.push(header.to_string()),
    }
    pieces
}

/// How many chunk matches to over-fetch per requested record in [`query`], so
/// de-duping chunks back to their parents still yields up to `k` distinct
/// records.
const OVERFETCH: usize = 8;

/// Separates a parent id from a chunk ordinal in a derived index key. ASCII Unit
/// Separator (never appears in real ids, which may themselves contain `#`/`/`),
/// so [`parent_key`] recovers the parent unambiguously.
const CHUNK_SEP: char = '\u{1f}';

/// The index key for chunk `ordinal` of the record at `parent`: same
/// namespace/collection, with the ordinal folded into the `id` (so a
/// [`KeyPrefix`] filter, which matches on namespace/collection only, is
/// unaffected). Every chunk — including a single-chunk fallback — carries an
/// ordinal, so keys are always parseable.
fn chunk_key(parent: &gonzalo_core::RecordKey, ordinal: usize) -> gonzalo_core::RecordKey {
    gonzalo_core::RecordKey::new(
        &parent.namespace,
        &parent.collection,
        format!("{}{CHUNK_SEP}{ordinal}", parent.id),
    )
}

/// Inverse of [`chunk_key`]: recover the parent key and chunk ordinal. A key
/// without the separator (defensive) is treated as parent, ordinal 0.
fn parent_key(chunk: &gonzalo_core::RecordKey) -> (gonzalo_core::RecordKey, usize) {
    match chunk.id.rsplit_once(CHUNK_SEP) {
        Some((id, ord)) => {
            let ordinal = ord.parse().unwrap_or(0);
            (
                gonzalo_core::RecordKey::new(&chunk.namespace, &chunk.collection, id),
                ordinal,
            )
        }
        None => (chunk.clone(), 0),
    }
}

/// Split text on blank-line boundaries into trimmed, non-empty paragraphs.
fn split_paragraphs(text: &str) -> Vec<String> {
    text.split("\n\n")
        .map(str::trim)
        .filter(|p| !p.is_empty())
        .map(str::to_string)
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use gonzalo_core::{Body, Identity, Meta, PutResult, RecordKey, Revision};
    use gonzalo_domain::{BodyFormat, Provider, State, StateCategory, TicketBody};
    use gonzalo_store_fs::FsStore;
    use gonzalo_vector::MemoryVectorIndex;
    use std::collections::BTreeMap;

    /// A deterministic bag-of-words embedder: cosine similarity tracks word
    /// overlap, enough to rank a matching record first.
    struct Bow;
    #[async_trait]
    impl Embedder for Bow {
        async fn embed(&self, text: &str) -> Result<Vec<f32>> {
            let mut v = vec![0f32; 32];
            for word in text.split_whitespace() {
                let h = word.bytes().map(|b| b as usize).sum::<usize>() % 32;
                v[h] += 1.0;
            }
            Ok(v)
        }
    }

    fn record(key: &RecordKey, kind: RecordKind, body: Body) -> Record {
        Record {
            revision: Revision::initial(body.bytes()),
            parent: None,
            body,
            kind,
            meta: Meta {
                author: Identity::new("t"),
                origin_system: "test".into(),
                created: 0,
                updated: 0,
                labels: BTreeMap::new(),
            },
            links: Vec::new(),
            key: key.clone(),
            ancestors: Vec::new(),
            deleted_at: None,
            deleted_blob: None,
            deleted_kind: None,
        }
    }

    #[test]
    fn chunk_splits_session_by_turn() {
        use gonzalo_domain::Turn;
        let session = Session {
            name: "sess".into(),
            turns: vec![
                Turn {
                    role: "user".into(),
                    text: "first turn about rust".into(),
                },
                Turn {
                    role: "assistant".into(),
                    text: "second turn about cooking".into(),
                },
                Turn {
                    role: "user".into(),
                    text: "third turn about music".into(),
                },
            ],
        };
        let key = RecordKey::new("caliban", "sessions", "s1");
        let chunks = chunk(&record(
            &key,
            RecordKind::Session,
            session.to_body().unwrap(),
        ))
        .unwrap()
        .unwrap();
        assert_eq!(chunks.len(), 3);
        // Session name is prepended to the first chunk so it stays searchable.
        assert!(chunks[0].contains("sess"));
        assert!(chunks[0].contains("first turn about rust"));
        assert!(chunks[1].contains("second turn about cooking"));
        assert!(chunks[2].contains("third turn about music"));
    }

    #[test]
    fn chunk_splits_topic_by_bullet_and_falls_back_to_one() {
        // Multiple bullets -> one chunk each.
        let topic = Topic {
            slug: "rust".into(),
            bullets: vec!["use clippy".into(), "run rustfmt".into()],
        };
        let key = RecordKey::new("caliban", "topics", "rust");
        let chunks = chunk(&record(&key, RecordKind::Topic, topic.to_body().unwrap()))
            .unwrap()
            .unwrap();
        assert_eq!(chunks.len(), 2);
        assert!(chunks[0].contains("rust"));
        assert!(chunks[0].contains("use clippy"));
        assert!(chunks[1].contains("run rustfmt"));

        // A single-piece record is the one-document fallback.
        let one = Topic {
            slug: "solo".into(),
            bullets: vec!["only bullet".into()],
        };
        let chunks = chunk(&record(&key, RecordKind::Topic, one.to_body().unwrap()))
            .unwrap()
            .unwrap();
        assert_eq!(chunks.len(), 1);

        // Non-knowledge kinds still gate to None.
        let ck = record(&key, RecordKind::Checkpoint, Body::Inline(b"{}".to_vec()));
        assert_eq!(chunk(&ck).unwrap(), None);
    }

    #[test]
    fn knowledge_text_per_kind() {
        let topic = Topic {
            slug: "rust".into(),
            bullets: vec!["use clippy".into()],
        };
        let key = RecordKey::new("caliban", "topics", "rust");
        let text = knowledge_text(&record(&key, RecordKind::Topic, topic.to_body().unwrap()));
        assert_eq!(text.as_deref(), Some("rust\nuse clippy"));

        // Checkpoint is not knowledge-bearing.
        let ck = record(&key, RecordKind::Checkpoint, Body::Inline(b"{}".to_vec()));
        assert_eq!(knowledge_text(&ck), None);
    }

    #[test]
    fn ticket_text_includes_title_body_labels() {
        let t = Ticket {
            provider: Provider::GitHub,
            uid: "o/r#1".into(),
            display: "#1".into(),
            item_type: "issue".into(),
            title: "fix the parser".into(),
            state: State {
                category: StateCategory::Open,
                resolution: None,
                raw_name: "open".into(),
                raw_id: None,
            },
            priority: None,
            actors: vec![],
            labels: vec!["bug".into()],
            containers: vec![],
            links: vec![],
            body: TicketBody {
                markdown: "the parser panics".into(),
                format: BodyFormat::Markdown,
                raw: None,
            },
            fields: BTreeMap::new(),
        };
        let key = RecordKey::new("tickets", "github", "o/r#1");
        let text = knowledge_text(&record(&key, RecordKind::Ticket, t.to_body().unwrap())).unwrap();
        assert!(text.contains("fix the parser"));
        assert!(text.contains("the parser panics"));
        assert!(text.contains("bug"));
    }

    async fn put(store: &FsStore, rec: Record) {
        assert!(matches!(
            store.put(rec, None).await.unwrap(),
            PutResult::Committed(_)
        ));
    }

    // --- gonzalo#326: ingest batches its index writes ---

    /// Counts how `KnowledgeStore` actually drives the index, so a regression to
    /// per-chunk upserts fails here rather than only showing up as extra
    /// manifest revisions on a durable backend.
    #[derive(Default)]
    struct CountingIndex {
        inner: MemoryVectorIndex,
        upsert: std::sync::atomic::AtomicUsize,
        upsert_many: std::sync::atomic::AtomicUsize,
        remove: std::sync::atomic::AtomicUsize,
    }

    impl CountingIndex {
        /// `(upsert, upsert_many, remove)` call counts.
        fn counts(&self) -> (usize, usize, usize) {
            use std::sync::atomic::Ordering::SeqCst;
            (
                self.upsert.load(SeqCst),
                self.upsert_many.load(SeqCst),
                self.remove.load(SeqCst),
            )
        }
    }

    #[async_trait]
    impl gonzalo_vector::VectorIndex for CountingIndex {
        async fn upsert(&self, key: RecordKey, vector: Vec<f32>) -> Result<()> {
            self.upsert
                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            self.inner.upsert(key, vector).await
        }
        async fn remove(&self, key: &RecordKey) -> Result<()> {
            self.remove
                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            self.inner.remove(key).await
        }
        async fn query(
            &self,
            query: &[f32],
            k: usize,
            filter: &KeyPrefix,
        ) -> Result<Vec<gonzalo_vector::Match>> {
            self.inner.query(query, k, filter).await
        }
        async fn keys(&self, filter: &KeyPrefix) -> Result<Vec<RecordKey>> {
            self.inner.keys(filter).await
        }
        async fn upsert_many(&self, items: Vec<(RecordKey, Vec<f32>)>) -> Result<()> {
            self.upsert_many
                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            self.inner.upsert_many(items).await
        }
    }

    fn topic_record(key: &RecordKey, bullets: &[&str]) -> Record {
        record(
            key,
            RecordKind::Topic,
            Topic {
                slug: key.id.clone(),
                bullets: bullets.iter().map(|b| (*b).to_string()).collect(),
            }
            .to_body()
            .unwrap(),
        )
    }

    /// A three-chunk record is one batched upsert, not three.
    #[tokio::test]
    async fn ingest_batches_a_records_chunks_into_one_upsert() {
        let dir = tempfile::tempdir().unwrap();
        let store = FsStore::new(dir.path());
        let key = RecordKey::new("caliban", "topics", "rust");
        put(
            &store,
            topic_record(&key, &["alpha one", "beta two", "gamma three"]),
        )
        .await;

        let index = std::sync::Arc::new(CountingIndex::default());
        let ks = KnowledgeStore::new(store, std::sync::Arc::clone(&index), Bow);
        assert!(ks.ingest(&key).await.unwrap());

        assert_eq!(
            index.counts(),
            (0, 1, 0),
            "one batch for the whole record, no per-chunk upsert"
        );
        assert_eq!(index.keys(&KeyPrefix::default()).await.unwrap().len(), 3);
    }

    /// A re-ingest after a shrink is still one batched upsert. The orphan
    /// removals remain one call each — see `ingest`'s note on why batching them
    /// waits on a workspace version bump — so this pins the current shape rather
    /// than claiming they are batched.
    #[tokio::test]
    async fn re_ingest_after_a_shrink_batches_the_upserts_and_drops_the_orphans() {
        let dir = tempfile::tempdir().unwrap();
        let store = FsStore::new(dir.path());
        let key = RecordKey::new("caliban", "topics", "rust");
        put(
            &store,
            topic_record(&key, &["alpha one", "beta two", "gamma three"]),
        )
        .await;

        let index = std::sync::Arc::new(CountingIndex::default());
        let ks = KnowledgeStore::new(FsStore::new(dir.path()), std::sync::Arc::clone(&index), Bow);
        assert!(ks.ingest(&key).await.unwrap());

        // Shrink 3 chunks -> 1, orphaning ordinals 1 and 2.
        let current = store.get(&key).await.unwrap().unwrap();
        assert!(matches!(
            store
                .put(topic_record(&key, &["alpha one"]), Some(current.revision))
                .await
                .unwrap(),
            PutResult::Committed(_)
        ));
        assert!(ks.ingest(&key).await.unwrap());

        let (upsert, upsert_many, remove) = index.counts();
        assert_eq!(upsert_many, 2, "one batch per ingest");
        assert_eq!(upsert, 0, "no per-chunk upsert");
        assert_eq!(remove, 2, "the two orphaned ordinals");
        assert_eq!(
            index.keys(&KeyPrefix::default()).await.unwrap().len(),
            1,
            "only the surviving chunk is indexed"
        );
    }

    /// The claim #326 actually makes, measured on the durable backend: ingesting
    /// one three-chunk record advances the vector manifest once, not three times.
    #[tokio::test]
    async fn ingest_commits_a_durable_manifest_once_per_record() {
        use gonzalo_vector::RecordVectorIndex;

        let dir = tempfile::tempdir().unwrap();
        let store = FsStore::new(dir.path());
        let key = RecordKey::new("caliban", "topics", "rust");
        put(
            &store,
            topic_record(&key, &["alpha one", "beta two", "gamma three"]),
        )
        .await;

        let manifest_key = gonzalo_core::VectorManifest::key("caliban", "topics");
        let index =
            RecordVectorIndex::open(FsStore::new(dir.path()), manifest_key.clone(), "bow-32", 32)
                .await
                .unwrap();
        let ks = KnowledgeStore::new(FsStore::new(dir.path()), index, Bow);

        assert!(ks.ingest(&key).await.unwrap());

        let manifest = store.get(&manifest_key).await.unwrap().expect("manifest");
        assert_eq!(
            manifest.revision.counter, 0,
            "one commit for a three-chunk record: the first commit on a fresh \
             manifest lands at counter 0 (`Revision::initial`), so a per-chunk \
             loop would show 2 here"
        );
    }

    #[tokio::test]
    async fn ingest_then_query_returns_matching_record() {
        let dir = tempfile::tempdir().unwrap();
        let store = FsStore::new(dir.path());

        let rust = RecordKey::new("caliban", "topics", "rust");
        let cooking = RecordKey::new("caliban", "topics", "cooking");
        put(
            &store,
            record(
                &rust,
                RecordKind::Topic,
                Topic {
                    slug: "rust".into(),
                    bullets: vec!["use clippy and cargo".into()],
                }
                .to_body()
                .unwrap(),
            ),
        )
        .await;
        put(
            &store,
            record(
                &cooking,
                RecordKind::Topic,
                Topic {
                    slug: "cooking".into(),
                    bullets: vec!["simmer the sauce slowly".into()],
                }
                .to_body()
                .unwrap(),
            ),
        )
        .await;

        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
        assert!(ks.ingest(&rust).await.unwrap());
        assert!(ks.ingest(&cooking).await.unwrap());

        let hits = ks
            .query("clippy cargo", 1, &KeyPrefix::default())
            .await
            .unwrap();
        assert_eq!(hits.len(), 1);
        assert_eq!(hits[0].record.key, rust);
    }

    fn turn(text: &str) -> gonzalo_domain::Turn {
        gonzalo_domain::Turn {
            role: "user".into(),
            text: text.into(),
        }
    }

    #[tokio::test]
    async fn query_ranks_a_matching_turn_over_a_diluted_decoy_and_dedups() {
        let dir = tempfile::tempdir().unwrap();
        let store = FsStore::new(dir.path());

        // A session whose one-document average is heavily diluted: two short
        // turns match "gamma", one long turn is all "alpha". Averaged into one
        // vector, the alpha turn drowns out the gamma signal.
        let sess = RecordKey::new("caliban", "sessions", "s1");
        put(
            &store,
            record(
                &sess,
                RecordKind::Session,
                Session {
                    name: "s".into(),
                    turns: vec![
                        turn("gamma"),
                        turn("gamma"),
                        turn("alpha alpha alpha alpha alpha"),
                    ],
                }
                .to_body()
                .unwrap(),
            ),
        )
        .await;

        // A decoy whose whole (short) document is closer to "gamma" than the
        // session's diluted average — so under one-document embedding the decoy
        // outranks the session.
        let decoy = RecordKey::new("caliban", "topics", "decoy");
        put(
            &store,
            record(
                &decoy,
                RecordKind::Topic,
                Topic {
                    slug: "gamma".into(),
                    bullets: vec!["alpha".into()],
                }
                .to_body()
                .unwrap(),
            ),
        )
        .await;

        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
        assert!(ks.ingest(&sess).await.unwrap());
        assert!(ks.ingest(&decoy).await.unwrap());

        let hits = ks.query("gamma", 5, &KeyPrefix::default()).await.unwrap();

        // Chunking lets the matching turn score on its own, so the session ranks
        // first...
        assert_eq!(hits[0].record.key, sess, "matching turn should rank first");
        // ...and the two matching turns de-dup to a single hit for the record.
        assert_eq!(
            hits.iter().filter(|h| h.record.key == sess).count(),
            1,
            "session should appear exactly once"
        );
    }

    #[tokio::test]
    async fn query_chunks_returns_individual_chunks_with_ordinals() {
        let dir = tempfile::tempdir().unwrap();
        let store = FsStore::new(dir.path());
        let sess = RecordKey::new("caliban", "sessions", "s1");
        put(
            &store,
            record(
                &sess,
                RecordKind::Session,
                Session {
                    name: "s".into(),
                    turns: vec![turn("gamma"), turn("gamma"), turn("alpha")],
                }
                .to_body()
                .unwrap(),
            ),
        )
        .await;

        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
        assert!(ks.ingest(&sess).await.unwrap());

        // The two "gamma" turns are the two closest chunks — returned
        // individually, NOT de-duped to the record.
        let hits = ks
            .query_chunks("gamma", 2, &KeyPrefix::default())
            .await
            .unwrap();
        assert_eq!(hits.len(), 2);
        assert!(hits.iter().all(|h| h.record.key == sess));
        let ordinals: std::collections::BTreeSet<usize> = hits.iter().map(|h| h.ordinal).collect();
        assert_eq!(ordinals, std::collections::BTreeSet::from([0, 1]));
        assert!(hits.iter().all(|h| h.text.contains("gamma")));
    }

    #[tokio::test]
    async fn reingest_shrink_removes_orphan_chunks() {
        let dir = tempfile::tempdir().unwrap();
        let sess = RecordKey::new("caliban", "sessions", "s1");
        let ks = KnowledgeStore::new(FsStore::new(dir.path()), MemoryVectorIndex::default(), Bow);

        // v1: three turns -> three chunks.
        put(
            ks.store(),
            record(
                &sess,
                RecordKind::Session,
                Session {
                    name: "s".into(),
                    turns: vec![turn("gamma"), turn("delta"), turn("epsilon")],
                }
                .to_body()
                .unwrap(),
            ),
        )
        .await;
        assert!(ks.ingest(&sess).await.unwrap());

        // The third turn's content is indexed while it exists.
        let before = ks
            .query_chunks("epsilon", 10, &KeyPrefix::default())
            .await
            .unwrap();
        assert!(before.iter().any(|h| h.text.contains("epsilon")));

        // v2: shrink to a single turn. The record now has one chunk; the "delta"
        // and "epsilon" chunks are orphaned in the index.
        let current = ks.store().get(&sess).await.unwrap().unwrap();
        let body = Session {
            name: "s".into(),
            turns: vec![turn("gamma")],
        }
        .to_body()
        .unwrap();
        let mut v2 = record(&sess, RecordKind::Session, body.clone());
        v2.parent = Some(current.revision.clone());
        v2.revision = current.revision.next(body.bytes());
        assert!(matches!(
            ks.store().put(v2, Some(current.revision)).await.unwrap(),
            PutResult::Committed(_)
        ));
        assert!(ks.ingest(&sess).await.unwrap());

        // The orphaned "epsilon" (and "delta") chunks are gone from the index;
        // only the surviving "gamma" chunk remains.
        let after = ks
            .query_chunks("epsilon", 10, &KeyPrefix::default())
            .await
            .unwrap();
        assert!(
            !after.iter().any(|h| h.text.contains("epsilon")),
            "orphaned chunk should have been removed on re-ingest"
        );
        assert!(
            !after.iter().any(|h| h.text.contains("delta")),
            "orphaned chunk should have been removed on re-ingest"
        );
        assert!(
            after.iter().any(|h| h.text.contains("gamma")),
            "the surviving turn should still be indexed"
        );
    }

    #[tokio::test]
    async fn ingest_skips_non_knowledge_kinds() {
        let dir = tempfile::tempdir().unwrap();
        let store = FsStore::new(dir.path());
        let key = RecordKey::new("caliban", "checkpoints", "c1");
        put(
            &store,
            record(&key, RecordKind::Checkpoint, Body::Inline(b"{}".to_vec())),
        )
        .await;
        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
        assert!(!ks.ingest(&key).await.unwrap(), "checkpoint is not indexed");
    }

    #[tokio::test]
    async fn ingest_surfaces_body_parse_error_and_gates_opaque_kinds() {
        let dir = tempfile::tempdir().unwrap();
        let store = FsStore::new(dir.path());

        // A knowledge-bearing kind (Topic) whose body bytes are not a valid
        // Topic: a corrupt record is a real failure, not "not indexable" (#139).
        let corrupt = RecordKey::new("caliban", "topics", "corrupt");
        put(
            &store,
            record(
                &corrupt,
                RecordKind::Topic,
                Body::Inline(b"not a valid topic".to_vec()),
            ),
        )
        .await;

        // A definitionally-opaque kind still gates cleanly to Ok(false).
        let ck = RecordKey::new("caliban", "checkpoints", "c1");
        put(
            &store,
            record(
                &ck,
                RecordKind::Checkpoint,
                Body::Inline(b"opaque".to_vec()),
            ),
        )
        .await;

        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
        assert!(
            ks.ingest(&corrupt).await.is_err(),
            "a corrupt knowledge-bearing body must surface an error, not Ok(false)"
        );
        assert!(
            !ks.ingest(&ck).await.unwrap(),
            "a genuinely-opaque kind is non-indexable: Ok(false)"
        );
    }

    #[tokio::test]
    async fn remove_deindexes_a_records_chunk_vectors() {
        let dir = tempfile::tempdir().unwrap();
        let store = FsStore::new(dir.path());

        // Three topics all matching "gamma". One will be removed; the other two
        // are the live set.
        let removed = RecordKey::new("caliban", "topics", "removed");
        let live1 = RecordKey::new("caliban", "topics", "live1");
        let live2 = RecordKey::new("caliban", "topics", "live2");
        for key in [&removed, &live1, &live2] {
            put(
                &store,
                record(
                    key,
                    RecordKind::Topic,
                    Topic {
                        slug: "gamma".into(),
                        bullets: vec!["gamma".into()],
                    }
                    .to_body()
                    .unwrap(),
                ),
            )
            .await;
        }

        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
        for key in [&removed, &live1, &live2] {
            assert!(ks.ingest(key).await.unwrap());
        }

        // De-index the one record. Its chunk vectors must no longer match, so
        // the k=2 query is filled entirely by the live records — no shrinkage
        // from orphaned vectors, and the removed record never appears (#150).
        ks.remove(&removed).await.unwrap();

        let hits = ks.query("gamma", 2, &KeyPrefix::default()).await.unwrap();
        assert_eq!(hits.len(), 2, "k=2 should be filled by the live records");
        let keys: std::collections::BTreeSet<RecordKey> =
            hits.iter().map(|h| h.record.key.clone()).collect();
        assert!(
            !keys.contains(&removed),
            "removed record must not appear after de-indexing"
        );
        assert_eq!(
            keys,
            std::collections::BTreeSet::from([live1, live2]),
            "only the live records should remain"
        );
    }

    // Re-ingesting a shrunk record must drop the chunks it no longer has — even
    // when the counts were not built in this process. Before `open`, a restart
    // resets them to zero and the orphans survive (#150, via #323).
    #[tokio::test]
    async fn counts_rebuilt_from_a_durable_index_still_drop_orphans() {
        use std::sync::Arc;

        let dir = tempfile::tempdir().unwrap();
        let index = Arc::new(MemoryVectorIndex::default());
        let key = RecordKey::new("ns", "coll", "doc");

        // `chunk()` maps a Topic to one chunk per bullet, so bullet count is
        // chunk count (crates/gonzalo-knowledge/src/lib.rs:250).
        let three = Topic {
            slug: "doc".into(),
            bullets: vec!["alpha".into(), "beta".into(), "gamma".into()],
        };
        put(
            &FsStore::new(dir.path()),
            record(&key, RecordKind::Topic, three.to_body().unwrap()),
        )
        .await;

        // A first "process" ingests all three chunks.
        let first = KnowledgeStore::new(FsStore::new(dir.path()), Arc::clone(&index), Bow);
        assert!(first.ingest(&key).await.unwrap());
        assert_eq!(index.keys(&KeyPrefix::default()).await.unwrap().len(), 3);
        drop(first);

        // The record shrinks to one bullet.
        let one = Topic {
            slug: "doc".into(),
            bullets: vec!["alpha".into()],
        };
        let handle = FsStore::new(dir.path());
        let existing = handle.get(&key).await.unwrap().unwrap();
        let shrunk = existing.update(
            one.to_body().unwrap(),
            Meta::new(Identity::new("test"), "test"),
        );
        assert!(matches!(
            handle.put(shrunk, Some(existing.revision)).await.unwrap(),
            PutResult::Committed(_)
        ));

        // A second "process" opens over the same durable index and re-ingests.
        // Chunks 1 and 2 must go; with `new` instead of `open` they would stay.
        let second = KnowledgeStore::open(FsStore::new(dir.path()), Arc::clone(&index), Bow)
            .await
            .unwrap();
        assert!(second.ingest(&key).await.unwrap());

        assert_eq!(index.keys(&KeyPrefix::default()).await.unwrap().len(), 1);
    }

    // F9 (#323 fix-round): a gap in stored ordinals must not undercount.
    // `open` rebuilds each record's count as `max(ordinal + 1)` over its
    // chunks, not the number of keys present, so ordinals {0, 2} (1 absent)
    // must rebuild to 3. A gap like this is reachable in practice: `ingest`'s
    // orphan-removal loop (`for ordinal in chunks.len()..old`) removes
    // ordinals one at a time in ascending order, so a crash partway through
    // can leave a lower ordinal removed while a higher one survives.
    #[tokio::test]
    async fn open_rebuilds_the_count_across_a_non_contiguous_ordinal_gap() {
        use std::sync::Arc;

        let dir = tempfile::tempdir().unwrap();
        let index = Arc::new(MemoryVectorIndex::default());
        let parent = RecordKey::new("ns", "coll", "doc");

        // Ordinals {0, 2} present, 1 absent. `MemoryVectorIndex` enforces one
        // dimension across everything it holds, so these must match what
        // `Bow` (32) will later embed on re-ingest.
        index
            .upsert(chunk_key(&parent, 0), vec![1.0; 32])
            .await
            .unwrap();
        index
            .upsert(chunk_key(&parent, 2), vec![1.0; 32])
            .await
            .unwrap();

        put(
            &FsStore::new(dir.path()),
            record(
                &parent,
                RecordKind::Topic,
                Topic {
                    slug: "doc".into(),
                    bullets: vec!["only bullet".into()],
                }
                .to_body()
                .unwrap(),
            ),
        )
        .await;

        let ks = KnowledgeStore::open(FsStore::new(dir.path()), Arc::clone(&index), Bow)
            .await
            .unwrap();

        // The record now has one chunk (ordinal 0). Re-ingesting must remove
        // ordinals 1 and 2 as orphans -- reachable only if `open` rebuilt the
        // count as 3 (the highest surviving ordinal + 1). A naive count of 2
        // (the number of keys present) would leave ordinal 2's entry behind,
        // since the removal loop would only run `1..2`.
        assert!(ks.ingest(&parent).await.unwrap());

        assert_eq!(
            index.keys(&KeyPrefix::default()).await.unwrap().len(),
            1,
            "the gap-created ordinal-2 orphan must have been removed on re-ingest"
        );
    }

    #[cfg(feature = "graph")]
    #[tokio::test]
    async fn query_in_subgraph_restricts_to_the_neighborhood() {
        use gonzalo_graph::{GraphStore, InMemoryGraphStore, build_rust};

        let dir = tempfile::tempdir().unwrap();
        let store = FsStore::new(dir.path());

        // Three topic records keyed by symbol name, all mentioning "database",
        // so all three are semantically similar to the query.
        for id in ["root", "helper", "faraway"] {
            let key = RecordKey::new("code", "symbols", id);
            let topic = Topic {
                slug: id.into(),
                bullets: vec!["touches the database layer".into()],
            };
            put(
                &store,
                record(&key, RecordKind::Topic, topic.to_body().unwrap()),
            )
            .await;
        }
        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
        for id in ["root", "helper", "faraway"] {
            assert!(
                ks.ingest(&RecordKey::new("code", "symbols", id))
                    .await
                    .unwrap()
            );
        }

        // Call graph: root -> helper; faraway is unrelated. So root's
        // neighborhood is {root, helper}.
        let mut graph = InMemoryGraphStore::new();
        graph.insert(
            "lib.rs",
            build_rust("fn root() { helper(); }\nfn helper() {}\nfn faraway() {}"),
        );

        let hits = ks
            .query_in_subgraph("database", 10, &graph, "root")
            .await
            .unwrap();
        let ids: BTreeMap<String, ()> =
            hits.iter().map(|h| (h.record.key.id.clone(), ())).collect();
        assert!(ids.contains_key("root"));
        assert!(ids.contains_key("helper"));
        assert!(
            !ids.contains_key("faraway"),
            "faraway is semantically similar but outside root's neighborhood"
        );
    }
}