kebab/crates/kebab-search/tests/hybrid.rs

//! Hybrid integration tests — touch a real `LanceVectorStore` +
//! `SqliteStore` + `MockEmbedder`. These tests are `#[ignore]`-d and
//! AVX-gated; see `tests/common/mod.rs` for the policy rationale.
//!
//! Mock-retriever unit tests live alongside the implementation in
//! `crates/kb-search/src/hybrid.rs` (no Lance, no AVX needed) — the
//! tests here exercise the full plumbing with the real Lance store.

mod common;

use std::path::PathBuf;
use std::sync::Arc;

use common::{
    HybridEnv, TEST_LEX_INDEX_VERSION, TEST_VEC_INDEX_VERSION, id32, require_avx_or_panic,
};
use kebab_core::{MediaType, Retriever, SearchFilters, SearchHit, SearchMode, SearchQuery};
use kebab_search::{FusionPolicy, HybridRetriever};
use rusqlite::params;
use serde_json::json;

fn build_hybrid(env: &HybridEnv) -> HybridRetriever {
    let lex: Arc<dyn Retriever> = Arc::new(env.lexical_retriever());
    let vec: Arc<dyn Retriever> = Arc::new(env.vector_retriever());
    HybridRetriever::with_policy(lex, vec, FusionPolicy::Rrf { k_rrf: 60 }, 5)
}

/// Seed a tiny corpus that lets us prove hybrid recall ≥ each side
/// independently. Two chunks are lexical-only matches ("rust cargo");
/// two chunks are vector-only matches (their text doesn't contain
/// the query token but their embedding still scores nearby because
/// MockEmbedder's hash distributes over all chunks).
fn seed_disjoint_corpus(env: &HybridEnv) -> Vec<String> {
    // The lexical side will only match chunks that contain the query
    // tokens. The vector side will rank ALL chunks by embedding
    // similarity to the query — even ones whose text doesn't share
    // a token with the query.
    let chunks = [
        // (chunk_id, doc_id, path, text, headings)
        (
            id32("c1"),
            id32("d1"),
            "notes/rust1.md",
            "rust cargo macros",
            &["A"][..],
        ),
        (
            id32("c2"),
            id32("d2"),
            "notes/rust2.md",
            "rust traits and lifetimes",
            &["B"][..],
        ),
        (
            id32("c3"),
            id32("d3"),
            "notes/python.md",
            "python dataclasses tutorial",
            &["C"][..],
        ),
        (
            id32("c4"),
            id32("d4"),
            "notes/go.md",
            "go interfaces and channels",
            &["D"][..],
        ),
    ];
    let mut ids = Vec::new();
    for (cid, did, path, text, headings) in &chunks {
        env.seed_chunk(cid, did, path, text, headings, &[]);
        env.embed_and_upsert(cid, did, text, headings);
        ids.push(cid.clone());
    }
    ids
}

#[test]
#[ignore = "requires AVX-capable hardware (LanceDB)"]
fn hybrid_recall_disjoint_returns_union() {
    require_avx_or_panic();
    let env = HybridEnv::new();
    let _ids = seed_disjoint_corpus(&env);
    let h = build_hybrid(&env);

    let q = SearchQuery {
        text: "rust".to_string(),
        mode: SearchMode::Hybrid,
        k: 4,
        filters: SearchFilters::default(),
    };
    let hits = h.search(&q).unwrap();

    // The vector side will return up to 4 candidates regardless of
    // text overlap; the lexical side will return only the rust* ones.
    // Together the union must cover at least the lexical hits AND
    // include at least one non-lexical chunk if vector found one.
    assert!(!hits.is_empty(), "hybrid must return at least one hit");
    // Every hit's RetrievalDetail.method must be Hybrid.
    for h in &hits {
        assert_eq!(h.retrieval.method, SearchMode::Hybrid);
        // At least one of lex/vec_score must be Some.
        assert!(
            h.retrieval.lexical_score.is_some() || h.retrieval.vector_score.is_some(),
            "hybrid hit must carry at least one mode's score"
        );
    }
    // index_version composite token.
    let iv = h.index_version();
    assert!(iv.0.starts_with("hybrid:"));
    assert!(iv.0.contains(TEST_LEX_INDEX_VERSION));
    assert!(iv.0.contains(TEST_VEC_INDEX_VERSION));

    // Lexical-only chunks (c1, c2) MUST appear: they're the only ones
    // matching the FTS5 query, and the vector side over-fetches enough
    // to include them too.
    let ids: Vec<&str> = hits.iter().map(|h| h.chunk_id.0.as_str()).collect();
    assert!(ids.contains(&id32("c1").as_str()));
    assert!(ids.contains(&id32("c2").as_str()));
}

#[test]
#[ignore = "requires AVX-capable hardware (LanceDB)"]
fn hybrid_determinism_same_query_twice() {
    require_avx_or_panic();
    let env = HybridEnv::new();
    let _ = seed_disjoint_corpus(&env);
    let h = build_hybrid(&env);

    let q = SearchQuery {
        text: "rust".to_string(),
        mode: SearchMode::Hybrid,
        k: 4,
        filters: SearchFilters::default(),
    };
    let a = h.search(&q).unwrap();
    let b = h.search(&q).unwrap();
    assert_eq!(
        a, b,
        "identical query must yield byte-identical Vec<SearchHit>"
    );
}

#[test]
#[ignore = "requires AVX-capable hardware (LanceDB)"]
fn hybrid_snapshot_run_1() {
    require_avx_or_panic();
    let env = HybridEnv::new();
    let _ = seed_disjoint_corpus(&env);
    let h = build_hybrid(&env);

    let q = SearchQuery {
        text: "rust".to_string(),
        mode: SearchMode::Hybrid,
        k: 4,
        filters: SearchFilters::default(),
    };
    let hits = h.search(&q).unwrap();

    // Snapshot pins the structural shape:
    //   - chunk_id ordering
    //   - which side contributed (lexical_rank / vector_rank
    //     populated as Some/None)
    //   - that fusion_score is non-increasing
    //   - method = Hybrid for every hit
    let actual = json!(
        hits.iter()
            .map(|h: &SearchHit| json!({
                "chunk_id": h.chunk_id.0,
                "rank": h.rank,
                "method": h.retrieval.method,
                "lexical_rank": h.retrieval.lexical_rank,
                "vector_rank": h.retrieval.vector_rank,
                "lex_some": h.retrieval.lexical_score.is_some(),
                "vec_some": h.retrieval.vector_score.is_some(),
                "fusion_score_positive": h.retrieval.fusion_score > 0.0,
            }))
            .collect::<Vec<_>>()
    );

    let fixture = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("fixtures")
        .join("search")
        .join("hybrid")
        .join("run-1.json");

    if std::env::var_os("KEBAB_UPDATE_SNAPSHOTS").is_some() {
        std::fs::create_dir_all(fixture.parent().unwrap()).unwrap();
        std::fs::write(&fixture, serde_json::to_string_pretty(&actual).unwrap()).unwrap();
        eprintln!("[snapshot] regenerated {}", fixture.display());
        // Fail loudly so that accidentally setting KEBAB_UPDATE_SNAPSHOTS
        // in CI surfaces as a test failure rather than a silent
        // overwrite + green run. Same fail-loud-instead-of-silent-pass
        // philosophy as P3-2's `SNAPSHOT_HASH_BASELINE = 0` and P3-3's
        // placeholder fixture guards.
        panic!(
            "[snapshot] regenerated {}, re-run without KEBAB_UPDATE_SNAPSHOTS to verify pin",
            fixture.display()
        );
    }

    let expected: serde_json::Value =
        serde_json::from_str(&std::fs::read_to_string(&fixture).unwrap_or_else(|_| {
            panic!(
                "missing snapshot fixture at {}; run with \
                 KEBAB_UPDATE_SNAPSHOTS=1 to create",
                fixture.display()
            )
        }))
        .unwrap();

    // Refuse to silently "pass" against the committed placeholder. The
    // placeholder JSON carries a `_comment` field with regeneration
    // instructions; production fixtures (a captured list) do not.
    assert!(
        !expected.get("_comment").is_some(),
        "snapshot fixture is a placeholder — regenerate on AVX hardware then commit. \
         Path: {}. To regenerate: \
         `KEBAB_UPDATE_SNAPSHOTS=1 cargo test -p kb-search -- --ignored hybrid_snapshot`.",
        fixture.display()
    );

    assert_eq!(
        actual, expected,
        "hybrid snapshot drift; rerun with KEBAB_UPDATE_SNAPSHOTS=1 to regenerate"
    );

    // Independent guard: fusion scores must be non-increasing across
    // the result list (rrf is rank-biased, so this is the
    // semantically-correct ordering invariant).
    for w in hits.windows(2) {
        assert!(
            w[0].retrieval.fusion_score >= w[1].retrieval.fusion_score,
            "fusion scores not in descending order: {} then {}",
            w[0].retrieval.fusion_score,
            w[1].retrieval.fusion_score
        );
    }
}

/// p9-fb-36: vector post-filter must pass `media` through `filter_chunks`.
/// Seeding two docs (markdown + pdf) and filtering for pdf-only must
/// return only the pdf chunk, proving `LanceVectorStore::search` →
/// `SqliteStore::filter_chunks` correctly applies the media arm.
#[test]
#[ignore = "requires AVX-capable hardware (LanceDB)"]
fn vector_filter_by_media() {
    require_avx_or_panic();
    let env = HybridEnv::new();
    env.insert_doc_with_media("md1.md", "rust ownership", MediaType::Markdown);
    env.insert_doc_with_media("doc.pdf", "rust pdf body", MediaType::Pdf);

    let filters = SearchFilters {
        media: vec!["pdf".to_string()],
        ..Default::default()
    };
    let hits = env.run_vector_search("rust", &filters);
    assert_eq!(hits.len(), 1, "media filter must keep only pdf chunk");
    assert!(
        hits[0].doc_path.0.ends_with(".pdf"),
        "expected .pdf path, got: {}",
        hits[0].doc_path.0
    );
}

/// p9-fb-36: vector post-filter must pass `doc_id` through `filter_chunks`.
/// Seeding two docs with shared text, filtering by one doc_id must return
/// only chunks from that doc.
#[test]
#[ignore = "requires AVX-capable hardware (LanceDB)"]
fn vector_filter_by_doc_id() {
    require_avx_or_panic();
    let env = HybridEnv::new();
    let target = env.insert_doc("a.md", "shared knowledge");
    env.insert_doc("b.md", "shared knowledge");

    let filters = SearchFilters {
        doc_id: Some(target.clone()),
        ..Default::default()
    };
    let hits = env.run_vector_search("shared", &filters);
    assert!(
        !hits.is_empty(),
        "doc_id filter must return hits for the target doc"
    );
    assert!(
        hits.iter().all(|h| h.doc_id == target),
        "all hits must belong to the target doc_id"
    );
}

#[test]
#[ignore = "requires AVX-capable hardware (LanceDB)"]
fn vector_hit_carries_indexed_at() {
    // p9-fb-32: VectorRetriever must populate SearchHit.indexed_at from
    // documents.updated_at via the JOIN added to hydrate_chunks (mirrors
    // the lexical retriever's behavior — Task 5).
    use time::OffsetDateTime;
    use time::format_description::well_known::Rfc3339;

    require_avx_or_panic();
    let env = HybridEnv::new();
    let _ids = seed_disjoint_corpus(&env);

    // `seed_chunk` hardcodes updated_at='1970-01-01T00:00:00Z'; bump
    // every document's updated_at to wall-clock now so the assertion
    // against `now` is meaningful.
    let now = OffsetDateTime::now_utc();
    let now_rfc = now.format(&Rfc3339).expect("format now as rfc3339");
    {
        let conn = env.sqlite.read_conn();
        conn.execute("UPDATE documents SET updated_at = ?", params![now_rfc])
            .expect("bump documents.updated_at");
    }

    let r = env.vector_retriever();
    let hits = r
        .search(&SearchQuery {
            text: "rust".to_string(),
            mode: SearchMode::Vector,
            k: 5,
            filters: SearchFilters::default(),
        })
        .expect("vector search");
    let hit = hits.first().expect("at least one vector hit");
    let now2 = OffsetDateTime::now_utc();
    let delta = (now2 - hit.indexed_at).whole_seconds().abs();
    assert!(delta < 60, "indexed_at within ±60s of now, got {delta}s");
    // stale is a placeholder set by the retriever; the App layer overwrites.
    assert!(!hit.stale, "vector retriever must default stale=false");
}

#[test]
#[ignore = "requires AVX-capable hardware (LanceDB)"]
fn vector_hit_carries_source_id_and_trust_level() {
    // rag-provenance-label: VectorRetriever must populate source_id +
    // trust_level from the documents join in hydrate_chunks. seed helpers
    // write trust_level='primary' and rely on V014 source_id DEFAULT
    // 'default'; override source_id to a non-default value to prove the join.
    use kebab_core::TrustLevel;

    require_avx_or_panic();
    let env = HybridEnv::new();
    let _ids = seed_disjoint_corpus(&env);
    {
        let conn = env.sqlite.read_conn();
        conn.execute("UPDATE documents SET source_id = 'notes'", [])
            .expect("set source_id");
    }

    let r = env.vector_retriever();
    let hits = r
        .search(&SearchQuery {
            text: "rust".to_string(),
            mode: SearchMode::Vector,
            k: 5,
            filters: SearchFilters::default(),
        })
        .expect("vector search");
    let hit = hits.first().expect("at least one vector hit");
    assert_eq!(hit.source_id.as_deref(), Some("notes"));
    assert_eq!(hit.trust_level, Some(TrustLevel::Primary));
}