src/vector/store.rs
//! Materialized embedding store — the warehouse layer for semantic search.
//!
//! Vectors live in Iceberg so search at any historical `git_sha` is a
//! **warehouse read** (no re-embed, no git walk) — the performance point for
//! mega codebases. Three tables (see [`crate::warehouse::iceberg_schema`]):
//!
//! - `embeddings` — content-addressed, deduped: one row per distinct
//! `(content_hash, model_profile)`, the vector a native `list<float32>`.
//! - `embedding_snapshots` — identity of one index build (`repo`, `git_sha`,
//! `model_profile`, `chunker_hash`, `fileset_hash`, `complete`).
//! - `embedding_manifest` — per-occurrence rows (`content_hash`, `file`,
//! `span`) belonging to a snapshot; no vector, no source text.
//!
//! Write: [`index_snapshot`] — batched dedup (one scan of existing hashes →
//! diff → batch-append only the missing vectors), idempotent per snapshot
//! identity. Read: [`reconstruct`] / [`search`] — pull a snapshot's manifest,
//! JOIN embeddings, build a [`VectorIndex`] in memory.
use std::collections::{HashMap, HashSet};
use anyhow::{Context, Result, anyhow, bail, ensure};
use arrow::array::{
Array, BooleanArray, Float32Array, Float32Builder, Int32Array, ListArray, ListBuilder,
RecordBatch, StringArray, TimestampMicrosecondArray,
};
use arrow::datatypes::{DataType, FieldRef, Schema as ArrowSchema};
use chrono::Utc;
use futures::TryStreamExt;
use iceberg::Catalog;
use iceberg::arrow::schema_to_arrow_schema;
use iceberg::expr::Reference;
use iceberg::spec::Datum;
use std::sync::{Arc, OnceLock};
use uuid::Uuid;
use super::chunk::Chunk;
use super::{HnswIndex, HnswParams, VectorIndex};
use crate::cache_budget::{Budget, BudgetedCache};
use crate::warehouse::iceberg::{
IcebergWarehouse, TABLE_EMBEDDING_MANIFEST, TABLE_EMBEDDING_SNAPSHOTS, TABLE_EMBEDDINGS,
append_batch,
};
/// Bump when the on-disk chunk semantics change so a re-index is forced.
/// `v2`: Mímir v2 removed the `syn` item-aware `.rs` chunking — every file is
/// line-windowed now, so v1 chunk sets are a different membership.
const CHUNKER_VERSION: &str = "v2";
// ----- identity types --------------------------------------------------------
/// Everything that determines what `embed(text)` produces. Two runs with the
/// same `ModelProfile` produce identical vectors — so it's the dedup key
/// alongside `content_hash` (not the bare model name, which can silently
/// change meaning across weight/tokenizer revisions).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelProfile {
pub model_name: String,
pub weights_sha: String,
pub tokenizer_sha: String,
pub pooling: String,
pub normalize: bool,
pub dim: usize,
pub dtype: String,
}
impl ModelProfile {
/// Stable id = SHA-256 over a canonical field encoding. This is the
/// `model_profile` string stored in every table.
pub fn id(&self) -> String {
let canonical = format!(
"model={}\0weights={}\0tokenizer={}\0pooling={}\0normalize={}\0dim={}\0dtype={}",
self.model_name,
self.weights_sha,
self.tokenizer_sha,
self.pooling,
self.normalize,
self.dim,
self.dtype,
);
hex_sha256(canonical.as_bytes())
}
}
/// Hash of the chunker version + its options — changing chunking changes
/// membership, so it's part of the snapshot identity.
pub fn chunker_hash(opts: &super::chunk::ChunkOptions) -> String {
hex_sha256(
format!(
"{CHUNKER_VERSION}\0max_lines={}\0overlap={}",
opts.max_lines, opts.overlap
)
.as_bytes(),
)
}
/// Turns text into vectors. Implemented by the embedder backends
/// (`embed-tract` / `embed-ort`); the store depends only on this seam, so it
/// stays testable and decoupled from the heavy model. `Send + Sync` so an
/// embedder can be shared across threads (CPU batch parallelism) and held in
/// an async server's state.
pub trait Embedder: Send + Sync {
/// Identity of this embedder — its `model_profile` (the dedup key).
fn profile(&self) -> ModelProfile;
/// Embed each text into a `profile().dim`-length, L2-normalized vector.
/// Output order matches input order. Takes **borrowed** `&[&str]` so the hot
/// path never clones a chunk's source text just to hand it to the embedder
/// (finding #1); a real backend builds one `[B, L]` batched submit from these.
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
/// **How long would embedding `total_chunks` take on THIS box, with THIS
/// backend?** Measured on a micro-sample, never by embedding the corpus.
///
/// This exists because the answer is not predictable from hardware. The same
/// jina weights are 2.38 chunk/s on `dwarf`'s ort-CPU path and would be two
/// orders of magnitude faster on a discrete card — and the Radeon 680M that
/// *looks* like acceleration is a lateral move once the model spills to GTT
/// (see `select::Probe::amd_gpu_is_worth_it`). Estimating that from vendor +
/// VRAM is the kind of arithmetic that is confidently wrong. Timing it is not.
///
/// A default method on purpose: it is written entirely in terms of [`embed`],
/// so every backend — ort, tract, static-table — gets it with no code of its
/// own. Override only if a backend can answer more cheaply than by running.
///
/// `sample` must come from the REAL corpus. Chunk cost scales with token
/// count, so a sample of short lines would flatter every backend equally and
/// still produce a useless number.
///
/// The first forward is embedded and **thrown away**: it pays first-touch page
/// faults on freshly-mmapped weights, and that is a load cost, not a per-chunk
/// cost. Counting it would penalise whichever backend has the larger file —
/// which is precisely the one quantization is trying to reward.
fn estimate(&self, sample: &[&str], total_chunks: usize) -> Result<Eta> {
ensure!(
sample.len() >= 2,
"estimate needs >= 2 sample chunks (one is consumed as an untimed warm-up)"
);
let _ = self
.embed(&sample[..1])
.context("warm-up forward failed; the backend cannot embed at all")?;
let timed = &sample[1..];
let t = std::time::Instant::now();
self.embed(timed).context("timed sample forward failed")?;
let secs = t.elapsed().as_secs_f64();
// The clock is this method's business; the arithmetic and the
// too-fast-to-time refusal are [`eta_from_sample`]'s, so both can be
// asserted without racing a real forward against a real microsecond.
eta_from_sample(timed.len(), secs, total_chunks)
}
}
/// **The longest an index build may take before it needs an explicit opt-in.**
///
/// Ten minutes is a human-scale number, not a measured one: it is the boundary
/// between "wait for it" and "go and do something else". It exists because the
/// alternative on `dwarf` is a 6.3 h fp32 pass (measured: 19 members, 22799 s)
/// that nothing warned about and nothing declined — the forge simply began, and
/// was still going when the box died.
pub const EMBED_BUDGET: std::time::Duration = std::time::Duration::from_secs(600);
/// What [`Embedder::estimate`] measured: a rate, and what it implies for the corpus.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Eta {
/// Measured throughput, chunk/s, warm-up excluded.
pub rate: f64,
/// How many chunks were actually timed (the warm-up is not among them).
pub sampled: usize,
/// Corpus size the estimate was extrapolated to.
pub total_chunks: usize,
/// `total_chunks / rate`.
pub eta: std::time::Duration,
}
impl Eta {
/// Does this fit [`EMBED_BUDGET`]? The gate the forge consults before starting.
pub fn within_budget(&self) -> bool {
self.eta <= EMBED_BUDGET
}
/// Does this fit an explicitly resolved budget (see [`embed_budget`])?
pub fn fits(&self, budget: std::time::Duration) -> bool {
self.eta <= budget
}
/// One line for an operator who has to decide whether to raise the budget.
pub fn summary(&self) -> String {
let mins = self.eta.as_secs_f64() / 60.0;
format!(
"{:.2} chunk/s (timed {} chunks) → {} chunks in ~{:.1} min — {}",
self.rate,
self.sampled,
self.total_chunks,
mins,
if self.within_budget() {
"within budget".to_string()
} else {
format!("OVER the {} min budget", EMBED_BUDGET.as_secs() / 60)
}
)
}
}
/// **The shortest sample wall-clock [`eta_from_sample`] will divide by.**
///
/// A microsecond. Below it the measurement is the clock's resolution rather than
/// the backend's speed, and dividing by it produces a rate that is arithmetic
/// noise dressed as a benchmark.
pub const MIN_TIMEABLE_SECS: f64 = 1e-6;
/// **The arithmetic half of [`Embedder::estimate`]** — a rate and an ETA from a
/// timed sample, or the refusal when the sample was too fast to time.
///
/// Split out from `estimate` so the refusal branch can be ENTERED ON PURPOSE.
/// `estimate` measures with `Instant::now()`, so a test that wanted to assert the
/// guard had to make a real embedder finish a real forward in under a
/// microsecond — which is a property of how busy the box is, not of this code.
/// MEASURED 2026-09-10 on t14s under seven parallel cargo builds: the single
/// timed forward took 1.293 µs, the guard did not fire, and
/// `estimate_refuses_to_time_the_untimeable` went red having found no defect
/// whatsoever. It raced the host clock and lost.
///
/// So the guard lives here, where `secs` is a parameter: the test passes `0.0`
/// and asserts the branch, on any box, at any load. `estimate` keeps the timing,
/// which is the part that genuinely needs a clock.
pub fn eta_from_sample(timed: usize, secs: f64, total_chunks: usize) -> Result<Eta> {
ensure!(timed > 0, "no chunks were timed (the sample was all warm-up)");
// A sample that lands under the clock's resolution would divide into a
// nonsense rate. Refuse rather than report a fast lie.
ensure!(
secs > MIN_TIMEABLE_SECS,
"sample of {timed} chunks completed too fast to time; pass more chunks"
);
let rate = timed as f64 / secs;
Ok(Eta {
rate,
sampled: timed,
total_chunks,
eta: std::time::Duration::from_secs_f64(total_chunks as f64 / rate),
})
}
/// Override for [`EMBED_BUDGET`], read by [`embed_budget`].
///
/// `off` / `none` / `0` disables the gate entirely — the operator has seen the
/// number and accepted it. Otherwise a duration: `600`, `600s`, `10m`, `2h`.
pub const NORNIR_EMBED_BUDGET_ENV: &str = "NORNIR_EMBED_BUDGET";
/// The FLOOR on how many chunks the gate embeds to price the job (one is spent
/// as the untimed warm-up). The actual sample is [`gate_sample_size`], which
/// also scales with the core count so the measured rate is the PARALLEL rate
/// the sustained run gets.
const GATE_SAMPLE_CHUNKS: usize = 8;
/// The gate's ACTUAL sample size: at least [`GATE_SAMPLE_CHUNKS`], and at least
/// one timed chunk per core (`available_parallelism + 1`, the +1 being the
/// untimed warm-up).
///
/// Why per-core: the CPU backends fan batches across cores, so a sustained
/// run's rate is a PARALLEL rate — and a sample smaller than the core count
/// cannot exercise that parallelism. Live on t14s (2026-08-26 13:11) the old
/// 7-timed-chunk sample formed ONE single-threaded batch, measured 0.14
/// chunk/s, and refused all 20 repos; the same box embeds several times faster
/// with every core busy. A core-count sample fans out one chunk per core
/// (`embed_support::fanout_rows_cap`), so it embeds in roughly the wall-clock
/// of its slowest chunk — CHEAPER than the old serial sample, not dearer.
fn gate_sample_size() -> usize {
let cores = std::thread::available_parallelism()
.map(|x| x.get())
.unwrap_or(1);
GATE_SAMPLE_CHUNKS.max(cores + 1)
}
/// How many `embed` calls [`check_embed_budget`] itself makes: one untimed
/// warm-up plus one timed sample.
///
/// This is a real cost, not an accounting detail — the gate embeds
/// [`gate_sample_size`] chunks and throws the vectors away, and it does so
/// BEFORE the durability loop. Negligible against a corpus of thousands, but it
/// is visible to anything that counts embedder invocations, so it is named here
/// rather than left for the next reader to rediscover from a failing fixture.
pub(crate) const GATE_PROBE_CALLS: usize = 2;
/// The active budget, or `None` when the gate is switched off.
///
/// Unparseable values are a **hard error at the gate**, not a silent fallback to
/// the default: an operator who typed `NORNIR_EMBED_BUDGET=10min` meant to change
/// the policy, and quietly enforcing 10 minutes anyway would look like it worked.
pub fn embed_budget() -> Result<Option<std::time::Duration>> {
match std::env::var(NORNIR_EMBED_BUDGET_ENV) {
Err(_) => Ok(Some(EMBED_BUDGET)),
Ok(raw) => parse_embed_budget(&raw),
}
}
/// The parse behind [`embed_budget`], as a pure function of the string.
///
/// Split out 2026-08-16 so the policy can be TESTED. It was reachable only through
/// `std::env::var`, which is process-global while the test binary runs in parallel, so
/// `the_budget_env_parses_units_and_refuses_nonsense` re-implemented the match inside its
/// own body and asserted its own copy — a test that stayed green if this function were
/// deleted outright. Two writers of one rule, and only one of them shipped.
///
/// Fixing it by construction (LAW 5) rather than by serialising env access: there is now
/// ONE parser, the caller does nothing but read the variable, and the test calls the same
/// function the gate does.
///
/// `Ok(None)` = the operator switched the gate off. `Err` = they meant something and
/// mistyped it, which is never a silent fallback to the default.
pub fn parse_embed_budget(raw: &str) -> Result<Option<std::time::Duration>> {
let v = raw.trim().to_ascii_lowercase();
if matches!(v.as_str(), "off" | "none" | "no" | "0") {
return Ok(None);
}
let (num, mult) = match v.as_bytes().last() {
Some(b'h') => (&v[..v.len() - 1], 3600.0),
Some(b'm') => (&v[..v.len() - 1], 60.0),
Some(b's') => (&v[..v.len() - 1], 1.0),
_ => (v.as_str(), 1.0),
};
let n: f64 = num.trim().parse().map_err(|_| {
anyhow!(
"{NORNIR_EMBED_BUDGET_ENV}=`{raw}` is not a duration; \
use `off` to disable, or a value like `600`, `600s`, `10m`, `2h`"
)
})?;
ensure!(
n.is_finite() && n > 0.0,
"{NORNIR_EMBED_BUDGET_ENV}=`{raw}` must be positive (use `off` to disable)"
);
Ok(Some(std::time::Duration::from_secs_f64(n * mult)))
}
/// Pick `want` texts spread EVENLY across `texts`, not the first `want`.
///
/// Chunk cost scales with token count, and a corpus is not shuffled: the first
/// chunks of a workspace are one file's, often its header. Striding samples short
/// and long chunks in the proportion the real run will meet.
fn stride_sample<'a>(texts: &[&'a str], want: usize) -> Vec<&'a str> {
if texts.len() <= want {
return texts.to_vec();
}
let step = texts.len() as f64 / want as f64;
(0..want)
.map(|i| texts[((i as f64 * step) as usize).min(texts.len() - 1)])
.collect()
}
/// **The gate: price the job, and refuse it if it cannot finish in time.**
///
/// Called with the texts that are actually about to be embedded — the deduped
/// missing set, not the whole corpus — so the estimate is of the real work.
///
/// Returns `Ok(Some(eta))` when it priced the job and the job fits, `Ok(None)`
/// when the gate could not or need not run, and `Err` when the job is over budget.
///
/// **A job it cannot price is allowed through.** Two cases reach that: a corpus
/// too small to sample, and a backend so fast the sample lands under the clock's
/// resolution — and the second is the static-table path, which is exactly the
/// backend that should never be blocked. A backend that is genuinely broken also
/// lands here, and that is deliberate: it will fail on the very next line with its
/// own error, which is a better message than one invented by the gate.
///
/// ⚠ **The budget is PER CALL, not per workspace.** `monitor`'s populate loop
/// calls the batched writer with one-member slices, so a 19-member workspace is
/// 19 independent gate decisions and the wall-clock ceiling is 19 × the budget.
/// That is the correct behaviour for a resumable per-member loop — each member
/// keeps what it finished — but it is NOT a total-runtime guarantee, and anyone
/// reading `EMBED_BUDGET` as "the forge cannot spend more than ten minutes"
/// would be wrong.
pub fn check_embed_budget<E: Embedder + ?Sized>(
embedder: &E,
texts: &[&str],
) -> Result<Option<Eta>> {
let Some(budget) = embed_budget()? else {
return Ok(None);
};
if texts.len() < 2 {
return Ok(None);
}
let sample = stride_sample(texts, gate_sample_size());
let eta = match embedder.estimate(&sample, texts.len()) {
Ok(e) => e,
Err(_) => return Ok(None),
};
ensure!(
eta.fits(budget),
"embedding {} chunks would take ~{:.1} min, over the {:.0} min budget \
({:.2} chunk/s warm sustained rate, measured on {} sample chunks embedded \
in parallel; model load is excluded — it is paid once at startup, not per \
chunk).\n\
Either pick a faster model (NORNIR_EMBED_MODEL — `potion-base-8m` is a \
static table, no forward pass) or accept the cost with \
{NORNIR_EMBED_BUDGET_ENV}=off (or e.g. {NORNIR_EMBED_BUDGET_ENV}=8h).",
texts.len(),
eta.eta.as_secs_f64() / 60.0,
budget.as_secs_f64() / 60.0,
eta.rate,
eta.sampled,
);
Ok(Some(eta))
}
/// Identity + tally of one index build.
#[derive(Debug, Clone)]
pub struct SnapshotRef {
pub snapshot_id: Uuid,
pub repo: String,
pub git_sha: String,
pub model_profile: String,
pub chunker_hash: String,
pub fileset_hash: String,
pub occurrences: usize,
pub new_vectors: usize,
}
/// Inputs that identify an index build.
pub struct IndexParams<'a> {
pub workspace: &'a str,
pub repo: &'a str,
pub git_sha: &'a str,
pub branch: &'a str,
pub model: &'a ModelProfile,
pub chunker_hash: &'a str,
/// Was the whole corpus indexed (vs a partial subset)? Search defaults to
/// complete snapshots.
pub complete: bool,
}
/// Identifies the repo + revision being indexed (the embedder/chunker are
/// supplied separately to [`index_repo`]).
pub struct RepoRef<'a> {
pub workspace: &'a str,
pub repo: &'a str,
pub git_sha: &'a str,
pub branch: &'a str,
/// Whole corpus (vs a partial subset)?
pub complete: bool,
}
/// One reconstructed search location.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Occurrence {
pub content_hash: String,
pub file: String,
pub start_line: usize,
pub end_line: usize,
}
/// Corpus size at/above which the live dense-kNN path prefers the HNSW ANN index
/// over the exact flat scan. Below it the flat scan is already sub-ms AND exact,
/// so building a graph is pure overhead; above it the O(log n) descent wins. Same
/// spirit as [`znippy_zoomies::vann::I8_SCAN_THRESHOLD`] — a *speed* switch.
const HNSW_MIN_ROWS: usize = 4_096;
/// A [`VectorIndex`] rebuilt from the warehouse plus the id → occurrence map. The
/// flat index is the exact **oracle**; [`Self::ann`] lazily builds an HNSW ANN
/// index (memoized) for the live approximate dense-kNN path on large corpora.
pub struct Reconstructed {
pub index: VectorIndex,
pub snapshot_id: Uuid,
pub git_sha: String,
/// `occurrence id (== manifest ordinal)` → location.
pub by_id: HashMap<u64, Occurrence>,
/// Lazily-built HNSW over `index`'s rows, memoized across queries against
/// this (immutable) snapshot. Only built when the corpus is large enough AND
/// `NORNIR_VECTOR_HNSW` isn't disabled — see [`Self::ann`].
hnsw: OnceLock<HnswIndex>,
}
impl Reconstructed {
/// **Resident heap bytes of this snapshot, INCLUDING everything built
/// lazily inside it** — the int8 matrix in `index`, and the HNSW graph in
/// `hnsw` (which keeps a second copy of the whole float matrix, so a hot
/// entry is roughly twice the size of the one that was admitted).
///
/// This is what [`INDEX_CACHE`] charges, and it is re-read after every query
/// rather than only at insert. See [`CachedIndex`].
pub fn resident_bytes(&self) -> usize {
let by_id = self.by_id.capacity()
* (std::mem::size_of::<(u64, Occurrence)>() + 16)
+ self
.by_id
.values()
.map(|o| o.content_hash.capacity() + o.file.capacity())
.sum::<usize>();
let hnsw = self.hnsw.get().map(|h| h.resident_bytes()).unwrap_or(0);
std::mem::size_of::<Self>()
+ self.index.resident_bytes()
+ by_id
+ self.git_sha.capacity()
+ hnsw
}
/// The HNSW ANN index for this snapshot, or `None` when the live path should
/// stay on the exact flat scan (small corpus, or `NORNIR_VECTOR_HNSW=0`). The
/// graph is built once on first request and reused (the flat index never
/// mutates after reconstruction, so a memoized graph can't go stale).
fn ann(&self) -> Option<&HnswIndex> {
if !hnsw_enabled() || self.index.len() < HNSW_MIN_ROWS {
return None;
}
Some(
self.hnsw
.get_or_init(|| self.index.build_hnsw(HnswParams::default())),
)
}
}
// ----- write path ------------------------------------------------------------
/// Walk `root` for `*.rs` source files (skipping any `target/` dir), returning
/// `(repo-relative path, contents)` sorted by path. Convenience for indexers
/// and benchmarks.
pub fn collect_rust_sources(root: &std::path::Path) -> Vec<(String, String)> {
let mut out = Vec::new();
for entry in walkdir::WalkDir::new(root)
.into_iter()
.filter_map(|e| e.ok())
{
let p = entry.path();
if !p.is_file() || p.extension().and_then(|e| e.to_str()) != Some("rs") {
continue;
}
if p.components().any(|c| c.as_os_str() == "target") {
continue;
}
if let Ok(s) = std::fs::read_to_string(p) {
let rel = p
.strip_prefix(root)
.unwrap_or(p)
.to_string_lossy()
.into_owned();
out.push((rel, s));
}
}
out.sort();
out
}
/// Chunk every file and persist a snapshot, embedding **only** content not
/// already stored for this model. Convenience wrapper over [`index_snapshot`]
/// that owns chunking + the embedder call.
pub fn index_repo<E: Embedder + ?Sized>(
wh: &IcebergWarehouse,
r: &RepoRef,
files: &[(String, String)],
opts: &super::chunk::ChunkOptions,
embedder: &E,
) -> Result<SnapshotRef> {
// Chunk files in parallel (item-aware for .rs, line-window otherwise). Each
// worker takes a contiguous slice of `files`; results are concatenated in
// file order so the snapshot's chunk list is deterministic. Same
// `std::thread::scope` + `available_parallelism` pattern as the embedder and
// the flat-ANN. (Embedding — the dominant cost — is already multicore in
// `embed.rs`.)
let chunks: Vec<Chunk> = if files.is_empty() {
Vec::new()
} else {
// Balanced (LPT + batched) fork-join map: each file is chunked by one
// worker, units claimed heaviest-first by content byte size. Per-file
// `Vec<Chunk>` results return in original `files` order, then flatten —
// identical to the serial concatenation, so the chunk list stays
// deterministic. (`0` workers ⇒ one per core, like the old pattern.)
znippy_zoomies::gatling_forkjoin::gatling_map_balanced(
files,
0,
1,
|(_, content): &(String, String)| content.len() as u64,
|_, (path, content)| super::chunk::chunk_file(path, content, opts),
)
.into_iter()
.flatten()
.collect()
};
let model = embedder.profile();
let ch = chunker_hash(opts);
let p = IndexParams {
workspace: r.workspace,
repo: r.repo,
git_sha: r.git_sha,
branch: r.branch,
model: &model,
chunker_hash: &ch,
complete: r.complete,
};
// STREAMED through the same slice door as the batched multi-repo path (gap
// 3): the missing chunks are embedded `embed_slice_chunks()` rows at a time
// and each slice is appended before the next is embedded, so the vectors in
// flight are bounded by the slice — not by the corpus — and a kill mid-run
// resumes from what already landed (vectors are content-addressed).
index_snapshot_streamed(
wh,
&p,
&chunks,
embed_slice_chunks(),
|missing| {
// Price the WHOLE job once before starting it. `missing` is already
// deduped against the warehouse, so this is the real work, not the
// nominal corpus. Borrow each chunk's text — no per-chunk String clone.
let texts: Vec<&str> = missing.iter().map(|c| c.text.as_str()).collect();
check_embed_budget(embedder, &texts).map(|_| ())
},
|slice| {
let texts: Vec<&str> = slice.iter().map(|c| c.text.as_str()).collect();
embedder.embed(&texts)
},
)
}
/// Rows per embed call on the per-repo path — [`EMBED_SLICE_CHUNKS`] unless
/// `NORNIR_EMBED_SLICE_CHUNKS` says otherwise (a positive integer). This is the
/// row-side ceiling of gap 3; the byte-side ceiling is `--max-memory`
/// (`vector::memory`), which bounds the fan-out INSIDE one embed call.
pub fn embed_slice_chunks() -> usize {
std::env::var("NORNIR_EMBED_SLICE_CHUNKS")
.ok()
.and_then(|v| v.trim().parse::<usize>().ok())
.filter(|&n| n > 0)
.unwrap_or(EMBED_SLICE_CHUNKS)
}
/// Owned inputs for one member in a [`index_repos_batched`] call — the
/// batched multi-repo sibling of [`index_repo`]'s per-repo [`RepoRef`] + files.
pub struct RepoInput {
pub workspace: String,
pub repo: String,
pub git_sha: String,
pub branch: String,
/// Whole corpus (vs a partial subset)?
pub complete: bool,
/// `(repo-relative path, contents)` — e.g. from [`collect_rust_sources`].
pub files: Vec<(String, String)>,
}
/// **Flattened, GPU-batched multi-repo embed** — the monitor's populate embed step.
///
/// The per-repo [`index_repo`] embeds ONE member at a time: it chunks that member's
/// files (parallel), then calls the embedder on just that member's missing chunks.
/// Run in a per-member loop, 31 of 32 cores sit idle between members AND the single
/// GPU sees one small batch per member instead of a few big ones. This function
/// fixes both:
///
/// 1. **CPU-side chunking is flattened** into ONE gatling pool over the UNION of
/// every member's files (heaviest-first by byte size), exactly like
/// `symbols::scan_repos_flattened` — so 1 big member or 17 small ones both
/// saturate every core, and the results group back per member deterministically.
/// 2. **The GPU forward is batched ACROSS members**: the missing content of ALL
/// members is de-duplicated into ONE set and embedded in a SINGLE `embedder.embed`
/// call, so the 4090 processes big batches (its own `plan_batches` bucketing then
/// packs them) instead of one-repo-at-a-time dribs. Vectors are written once.
/// 3. **Only the per-member snapshot/manifest APPENDS stay serial** (single-writer,
/// zero compute) — the same principle as the symbol scan's serial appends.
///
/// Returns one `Result<SnapshotRef>` per input member, in input order — a member's
/// failure is isolated to its slot and never aborts the batch. Idempotent per
/// member: a member whose identical snapshot already exists is returned unchanged
/// with `new_vectors == 0` and contributes nothing to the embed.
/// How many chunks [`index_repos_batched`] hands the embedder per call.
///
/// **This does not make the embed faster, and is not claimed to.** Ingest is
/// EMBED-BOUND, and the embed is bound by the inference device, not by how the
/// caller slices its input: measured full-from-empty on **oden** (Threadripper
/// PRO 3975WX, 32 cores, RTX 4090) the embed is **165.1 s of a 170.5 s populate —
/// 96.9 % of wall — at `cores_busy = 1.12/32`**, because `ort`'s `Session`
/// serialises on the one GPU. There is no caller-side parallelism left to win;
/// the levers are a bigger forward batch, a second device, or a smaller model.
/// (`workspace_nordisk/.nornir/mimir-ingest-measurement-2026-07-26.md`.)
///
/// This constant is the **durability + observability slice**, and only that:
/// peak resident text+vectors is bounded by it, one Parquet append lands per
/// slice, and one progress line is printed per slice. It is NOT the forward-pass
/// batch size — every embedder does its own `plan_batches` bucketing inside
/// `embed`, so the device still sees the batches it wants.
///
/// What it buys, on the host where it matters: on **t14s** (16 logical / 8
/// physical hybrid cores, **no GPU**, so the embedder falls back to CPU
/// `tract-onnx`) the same populate's embed ran **>62 minutes without completing**,
/// at 7.3 GB RSS, having written **zero** vectors — one monolithic
/// `embedder.embed(&all_texts)`, so a kill or an OOM anywhere in that hour threw
/// away the whole hour, and nothing on stderr distinguished *working* from *hung*.
/// 2048 caps the all-or-nothing window at one slice and makes a retry resume from
/// what already landed (vectors are content-addressed). ≈20 extra appends for a
/// 40k-chunk workspace populate — a rounding error next to what `warehouse
/// compact` already handles.
const EMBED_SLICE_CHUNKS: usize = 2048;
pub fn index_repos_batched<E: Embedder + ?Sized>(
wh: &IcebergWarehouse,
inputs: &[RepoInput],
opts: &super::chunk::ChunkOptions,
embedder: &E,
) -> Vec<Result<SnapshotRef>> {
if inputs.is_empty() {
return Vec::new();
}
// Phase 1 — CPU (parallel): chunk the UNION of every member's files in ONE
// flattened gatling pool (heaviest-first by content byte size), then group the
// per-file chunk lists back per member in (member, file) order ⇒ deterministic,
// identical to chunking each member serially.
let mut union: Vec<(usize, &str, &str)> = Vec::new();
for (mi, inp) in inputs.iter().enumerate() {
for (path, content) in &inp.files {
union.push((mi, path.as_str(), content.as_str()));
}
}
let per_file: Vec<(usize, Vec<Chunk>)> = if union.is_empty() {
Vec::new()
} else {
znippy_zoomies::gatling_forkjoin::gatling_map_balanced(
&union,
0,
1,
|(_, _, content): &(usize, &str, &str)| content.len() as u64,
|_, (mi, path, content)| (*mi, super::chunk::chunk_file(path, content, opts)),
)
};
let mut chunks_by_member: Vec<Vec<Chunk>> = (0..inputs.len()).map(|_| Vec::new()).collect();
for (mi, mut cs) in per_file {
chunks_by_member[mi].append(&mut cs);
}
// The embed + all appends touch the warehouse ⇒ one async section.
let model = embedder.profile();
let ch = chunker_hash(opts);
let dim = model.dim;
let model_profile = model.id();
wh.block_on(async {
// Existing content is read ONCE for the whole batch (not per member).
let existing = match existing_content_hashes(wh, &model_profile).await {
Ok(e) => e,
Err(e) => {
// A read failure fails every member uniformly (nothing was written).
return (0..inputs.len())
.map(|_| Err(anyhow!("read existing content hashes: {e:#}")))
.collect::<Vec<_>>();
}
};
// Per-member identity + idempotent skip. `todo` = members that need a new
// snapshot; `out[i]` is pre-filled for skipped/failed members.
let mut out: Vec<Option<Result<SnapshotRef>>> = (0..inputs.len()).map(|_| None).collect();
let mut todo: Vec<usize> = Vec::new();
for (i, inp) in inputs.iter().enumerate() {
let chunks = &chunks_by_member[i];
let fileset_hash = fileset_hash_of(chunks);
match find_snapshot(
wh,
&inp.repo,
&inp.git_sha,
&model_profile,
&ch,
&fileset_hash,
)
.await
{
Ok(Some(existing_id)) => {
out[i] = Some(Ok(SnapshotRef {
snapshot_id: existing_id,
repo: inp.repo.clone(),
git_sha: inp.git_sha.clone(),
model_profile: model_profile.clone(),
chunker_hash: ch.clone(),
fileset_hash,
occurrences: chunks.len(),
new_vectors: 0,
}));
}
Ok(None) => todo.push(i),
Err(e) => out[i] = Some(Err(anyhow!("find snapshot for {}: {e:#}", inp.repo))),
}
}
// Phase 2 — the GLOBAL missing set across all `todo` members (first
// occurrence wins, order preserved), embedded in ONE batched call so the
// single GPU sees big batches instead of one-repo-at-a-time dribs.
// Borrow the chunks (no whole-`Chunk` clone) and intern the dedup keys as
// `&str` into `chunks_by_member` (no `content_hash.to_string()` per chunk)
// — finding #1's collection-layer clones.
let mut missing: Vec<&Chunk> = Vec::new();
let mut seen: HashSet<&str> = HashSet::new();
for &i in &todo {
for c in &chunks_by_member[i] {
let h = c.content_hash.as_str();
if !existing.contains(h) && seen.insert(h) {
missing.push(c);
}
}
}
// The set of content hashes newly embedded by THIS batch (for per-member
// `new_vectors` attribution below) — borrowed, not cloned.
let newly: HashSet<&str> = missing.iter().map(|c| c.content_hash.as_str()).collect();
// STREAMED in bounded slices, not one monolithic call. NOT a speedup —
// see [`EMBED_SLICE_CHUNKS`]; ingest is device-bound and slicing changes
// the throughput by nothing. It is a DURABILITY + OBSERVABILITY change.
//
// Measured on t14s (2026-07-27, CPU `tract-onnx`, no GPU): a FIRST
// populate of the 17-member nordisk workspace makes every chunk "missing",
// and the old single `embedder.embed(&all_texts)` then ran **>62 minutes
// without completing** while
// * holding every input text AND every output vector resident — the
// process reached **7.3 GB RSS** and grew monotonically;
// * writing NOTHING until the very end, so a kill/OOM at minute 61 left
// zero vectors on disk and the whole hour had to be redone (it was
// killed twice, and twice cost the entire run);
// * printing NOTHING between "embedding 17 member(s) …" and the end, so
// an operator could not distinguish *working* from *hung* — the
// Unknown-≠-Idle failure, one phase up.
// Slicing costs nothing (the embedder's own `plan_batches` bucketing runs
// per call and already caps the forward-pass batch) and buys bounded peak
// memory, per-slice durability, and a live rate/ETA line. That is what
// makes the ecosystem-scale question answerable at all: at 10³–10⁵ crates
// an all-or-nothing embed cannot finish, and a run you cannot observe is a
// run you cannot cost.
// ── THE BUDGET GATE ──────────────────────────────────────────────────
// Price the job BEFORE the first slice, not after the first hour. `missing`
// is already deduped against the warehouse, so it is the real work and not
// the nominal corpus — a re-populate that only has to embed 40 new chunks
// is priced as 40, and passes, even on the box where a cold run would not.
//
// A refusal here writes NOTHING: no slice has run, so unlike the mid-embed
// failure below there is no partial state to explain. Every `todo` member
// gets the same error, because none of their snapshots can be completed.
if !missing.is_empty() {
let all_texts: Vec<&str> = missing.iter().map(|c| c.text.as_str()).collect();
match check_embed_budget(embedder, &all_texts) {
Ok(Some(eta)) => eprintln!("nornir-embed: {}", eta.summary()),
// Unpriceable (too fast to time, or too small to sample) ⇒ proceed.
Ok(None) => {}
Err(e) => {
let msg = format!("{e:#}");
for &i in &todo {
out[i] = Some(Err(anyhow!("{msg}")));
}
return out.into_iter().map(|o| o.unwrap()).collect();
}
}
}
if !missing.is_empty() {
let total = missing.len();
let t_embed = std::time::Instant::now();
let mut done = 0usize;
let mut failed: Option<String> = None;
for slice in missing.chunks(EMBED_SLICE_CHUNKS) {
// Borrow each chunk's text for the batched embed — no String clone.
let texts: Vec<&str> = slice.iter().map(|c| c.text.as_str()).collect();
let vectors = match embedder.embed(&texts) {
Ok(v) => v,
Err(e) => {
failed = Some(format!("embed batch: {e:#}"));
break;
}
};
if vectors.len() != slice.len() {
failed = Some(format!(
"embedder returned {} vectors for {} chunks",
vectors.len(),
slice.len()
));
break;
}
let mut to_write: Vec<(&str, &[f32])> = Vec::with_capacity(slice.len());
let mut dim_err: Option<String> = None;
for (&c, v) in slice.iter().zip(&vectors) {
if let Err(e) = ensure_dim(v, dim, c) {
dim_err = Some(format!("{e:#}"));
break;
}
to_write.push((c.content_hash.as_str(), v.as_slice()));
}
if let Some(msg) = dim_err {
failed = Some(msg);
break;
}
// DURABLE PER SLICE: the vectors are content-addressed, so a run
// that dies here resumes from `existing_content_hashes` next time
// instead of restarting from zero.
if let Err(e) =
append_embeddings(wh, &model_profile, dim, model.normalize, &to_write).await
{
failed = Some(format!("append embeddings: {e:#}"));
break;
}
done += slice.len();
let secs = t_embed.elapsed().as_secs_f64();
let rate = if secs > 0.0 { done as f64 / secs } else { 0.0 };
let eta = if rate > 0.0 {
(total - done) as f64 / rate
} else {
f64::NAN
};
eprintln!(
"nornir-embed: {done}/{total} chunk(s) ({:.0}%) · {rate:.0} chunk/s · \
elapsed {secs:.0}s · eta {eta:.0}s",
100.0 * done as f64 / total as f64,
);
}
if let Some(msg) = failed {
// Slices already committed STAY committed (content-addressed, so
// they are correct and re-usable) — but every member is failed,
// because no member's snapshot is complete.
let msg = format!("{msg} (after {done}/{total} chunk(s) persisted)");
for &i in &todo {
out[i] = Some(Err(anyhow!("{msg}")));
}
return out.into_iter().map(|o| o.unwrap()).collect();
}
}
// Phase 3 — SERIAL single-writer appends: one manifest + snapshot row per
// `todo` member (vectors already written above). `credited` attributes each
// newly-embedded content hash to the FIRST `todo` member that carries it, so
// per-member `new_vectors` sums to the batch total exactly like the serial
// per-member path (where the first member to introduce content embeds it and
// later members see it already stored) — no cross-member double count.
let mut credited: HashSet<String> = HashSet::new();
for &i in &todo {
let inp = &inputs[i];
let chunks = &chunks_by_member[i];
let fileset_hash = fileset_hash_of(chunks);
let snapshot_id = Uuid::new_v4();
let p = IndexParams {
workspace: &inp.workspace,
repo: &inp.repo,
git_sha: &inp.git_sha,
branch: &inp.branch,
model: &model,
chunker_hash: &ch,
complete: inp.complete,
};
let res = async {
append_manifest(wh, snapshot_id, &model_profile, chunks).await?;
append_snapshot_row(
wh,
snapshot_id,
&p,
&model_profile,
&fileset_hash,
chunks.len(),
)
.await?;
Ok::<(), anyhow::Error>(())
}
.await;
out[i] = Some(match res {
Ok(()) => {
let new_vectors = chunks
.iter()
.filter(|c| {
newly.contains(c.content_hash.as_str())
&& credited.insert(c.content_hash.clone())
})
.count();
Ok(SnapshotRef {
snapshot_id,
repo: inp.repo.clone(),
git_sha: inp.git_sha.clone(),
model_profile: model_profile.clone(),
chunker_hash: ch.clone(),
fileset_hash,
occurrences: chunks.len(),
new_vectors,
})
}
Err(e) => Err(anyhow!("persist snapshot for {}: {e:#}", inp.repo)),
});
}
out.into_iter().map(|o| o.unwrap()).collect()
})
}
/// Persist `chunks` (all occurrences, in order) as a new snapshot. For each
/// distinct `content_hash` **not** already stored for this `model_profile`,
/// `embed_missing` is invoked once — with the missing chunks in a stable
/// order — to produce their vectors. Unchanged content is never re-embedded
/// (the performance point). Idempotent: a snapshot with the same identity
/// (incl. `fileset_hash`) returns without writing or embedding.
pub fn index_snapshot<F>(
wh: &IcebergWarehouse,
p: &IndexParams,
chunks: &[Chunk],
embed_missing: F,
) -> Result<SnapshotRef>
where
F: FnOnce(&[Chunk]) -> Result<Vec<Vec<f32>>>,
{
// The one-call door: a `usize::MAX` slice makes the streamed path call the
// embedder exactly once with every missing chunk — the historical contract
// (`unchanged_content_dedups_across_commits` asserts on that single call).
let mut once = Some(embed_missing);
index_snapshot_streamed(wh, p, chunks, usize::MAX, |_| Ok(()), move |slice| {
(once
.take()
.expect("a usize::MAX slice yields exactly one embed call"))(slice)
})
}
/// [`index_snapshot`] with the embed STREAMED in slices of `slice_rows` missing
/// chunks, each slice appended to `embeddings` before the next is embedded.
///
/// `price` runs ONCE over the whole missing set before any embedding (the
/// budget gate's place); `embed_slice` runs per slice and is `FnMut`. The
/// vectors in flight are bounded by `slice_rows` × dim, and a run that dies
/// mid-way resumes from `existing_content_hashes` next time — the same
/// durability the batched multi-repo path already had, now on the per-repo
/// `vector index <repo>` path (gap 3 of INVENTORY-semantic-twins §5).
pub fn index_snapshot_streamed<P, F>(
wh: &IcebergWarehouse,
p: &IndexParams,
chunks: &[Chunk],
slice_rows: usize,
price: P,
mut embed_slice: F,
) -> Result<SnapshotRef>
where
P: FnOnce(&[Chunk]) -> Result<()>,
F: FnMut(&[Chunk]) -> Result<Vec<Vec<f32>>>,
{
let dim = p.model.dim;
let model_profile = p.model.id();
let fileset_hash = fileset_hash_of(chunks);
wh.block_on(async {
if let Some(existing) = find_snapshot(
wh,
p.repo,
p.git_sha,
&model_profile,
p.chunker_hash,
&fileset_hash,
)
.await?
{
return Ok(SnapshotRef {
snapshot_id: existing,
repo: p.repo.into(),
git_sha: p.git_sha.into(),
model_profile,
chunker_hash: p.chunker_hash.into(),
fileset_hash,
occurrences: chunks.len(),
new_vectors: 0,
});
}
// Which content_hashes for this model are already stored?
let existing = existing_content_hashes(wh, &model_profile).await?;
// Distinct chunks whose content isn't stored yet (first occurrence
// wins, order preserved) — only these get embedded.
let mut missing: Vec<Chunk> = Vec::new();
let mut seen: HashSet<&str> = HashSet::new();
for c in chunks {
let h = c.content_hash.as_str();
if !existing.contains(h) && seen.insert(h) {
missing.push(c.clone());
}
}
let new_vectors = missing.len();
if !missing.is_empty() {
price(&missing)?;
let slice_rows = slice_rows.max(1);
let total = missing.len();
let streamed = slice_rows < total;
let t_embed = std::time::Instant::now();
let mut done = 0usize;
for slice in missing.chunks(slice_rows) {
let vectors = embed_slice(slice)?;
if vectors.len() != slice.len() {
bail!(
"embedder returned {} vectors for {} chunks",
vectors.len(),
slice.len()
);
}
// Sized to the slice, not the corpus: this is the bound.
let mut to_write: Vec<(&str, &[f32])> = Vec::with_capacity(slice.len());
for (c, v) in slice.iter().zip(&vectors) {
ensure_dim(v, dim, c)?;
to_write.push((c.content_hash.as_str(), v.as_slice()));
}
// DURABLE PER SLICE — a kill here loses this slice, not the run.
append_embeddings(wh, &model_profile, dim, p.model.normalize, &to_write).await?;
done += slice.len();
if streamed {
let secs = t_embed.elapsed().as_secs_f64();
let rate = if secs > 0.0 { done as f64 / secs } else { 0.0 };
let eta = if rate > 0.0 { (total - done) as f64 / rate } else { f64::NAN };
eprintln!(
"nornir-embed: {done}/{total} chunk(s) ({:.0}%) · {rate:.0} chunk/s · \
elapsed {secs:.0}s · eta {eta:.0}s",
100.0 * done as f64 / total as f64,
);
}
}
}
let snapshot_id = Uuid::new_v4();
append_manifest(wh, snapshot_id, &model_profile, chunks).await?;
append_snapshot_row(
wh,
snapshot_id,
p,
&model_profile,
&fileset_hash,
chunks.len(),
)
.await?;
Ok(SnapshotRef {
snapshot_id,
repo: p.repo.into(),
git_sha: p.git_sha.into(),
model_profile,
chunker_hash: p.chunker_hash.into(),
fileset_hash,
occurrences: chunks.len(),
new_vectors,
})
})
}
// ----- read path -------------------------------------------------------------
/// Row counts across the three vector tables.
#[derive(Debug, Default, Clone)]
pub struct WarehouseStats {
pub embeddings: usize,
pub snapshots: usize,
pub occurrences: usize,
}
/// Count rows in `embeddings`, `embedding_snapshots`, `embedding_manifest`.
pub fn warehouse_stats(wh: &IcebergWarehouse) -> Result<WarehouseStats> {
wh.block_on(async {
Ok(WarehouseStats {
embeddings: count_rows(wh, TABLE_EMBEDDINGS).await?,
snapshots: count_rows(wh, TABLE_EMBEDDING_SNAPSHOTS).await?,
occurrences: count_rows(wh, TABLE_EMBEDDING_MANIFEST).await?,
})
})
}
async fn count_rows(wh: &IcebergWarehouse, table: &str) -> Result<usize> {
let t = wh.catalog().load_table(&wh.table_ident(table)).await?;
let scan = t.scan().select(["model_profile"]).build()?;
let batches: Vec<RecordBatch> = scan.to_arrow().await?.try_collect().await?;
Ok(batches.iter().map(|b| b.num_rows()).sum())
}
/// Rebuild the [`VectorIndex`] for `(repo, model_profile)` at `git_sha`
/// (latest snapshot when `git_sha` is `None`). Pure warehouse read.
pub fn reconstruct(
wh: &IcebergWarehouse,
repo: &str,
git_sha: Option<&str>,
model_profile: &str,
) -> Result<Reconstructed> {
wh.block_on(async {
let (snapshot_id, sha) = resolve_snapshot(wh, repo, git_sha, model_profile).await?;
build_reconstructed(wh, snapshot_id, sha, model_profile).await
})
}
/// Read the raw `(occurrence, vector)` rows of `(repo, model_profile)`'s latest
/// (or `git_sha`-pinned) snapshot, in manifest-ordinal order — the flat data a
/// lens FOLDS directly (e.g. SemanticClones consuming `<repo>::slices`),
/// without building a search index. Same resolution + join as [`reconstruct`];
/// errors identically when the namespace was never indexed.
pub fn reconstruct_rows(
wh: &IcebergWarehouse,
repo: &str,
git_sha: Option<&str>,
model_profile: &str,
) -> Result<Vec<(Occurrence, Vec<f32>)>> {
wh.block_on(async {
let (snapshot_id, _sha) = resolve_snapshot(wh, repo, git_sha, model_profile).await?;
let manifest = read_manifest(wh, snapshot_id).await?;
if manifest.is_empty() {
bail!("snapshot {snapshot_id} has an empty manifest");
}
let needed: HashSet<&str> = manifest.iter().map(|m| m.content_hash.as_str()).collect();
let (vectors, _dim) = read_vectors(wh, model_profile, &needed).await?;
let mut out = Vec::with_capacity(manifest.len());
for occ in &manifest {
let v = vectors.get(occ.content_hash.as_str()).ok_or_else(|| {
anyhow!("manifest references missing vector {}", occ.content_hash)
})?;
out.push((
Occurrence {
content_hash: occ.content_hash.clone(),
file: occ.file.clone(),
start_line: occ.start_line,
end_line: occ.end_line,
},
v.clone(),
));
}
Ok(out)
})
}
/// Build a [`Reconstructed`] for an already-resolved snapshot (manifest JOIN
/// embeddings → in-memory index). The expensive step that the cache memoizes.
async fn build_reconstructed(
wh: &IcebergWarehouse,
snapshot_id: Uuid,
sha: String,
model_profile: &str,
) -> Result<Reconstructed> {
let manifest = read_manifest(wh, snapshot_id).await?;
if manifest.is_empty() {
bail!("snapshot {snapshot_id} has an empty manifest");
}
let needed: HashSet<&str> = manifest.iter().map(|m| m.content_hash.as_str()).collect();
let (vectors, dim) = read_vectors(wh, model_profile, &needed).await?;
let mut index = VectorIndex::new(dim)?;
let mut by_id = HashMap::with_capacity(manifest.len());
// One vector per occurrence (id == ordinal); identical content shares
// the same vector under distinct ids.
let mut flat: Vec<f32> = Vec::with_capacity(manifest.len() * dim);
let mut ids: Vec<u64> = Vec::with_capacity(manifest.len());
for occ in &manifest {
let v = vectors
.get(occ.content_hash.as_str())
.ok_or_else(|| anyhow!("manifest references missing vector {}", occ.content_hash))?;
flat.extend_from_slice(v);
ids.push(occ.ordinal as u64);
by_id.insert(
occ.ordinal as u64,
Occurrence {
content_hash: occ.content_hash.clone(),
file: occ.file.clone(),
start_line: occ.start_line,
end_line: occ.end_line,
},
);
}
index.add(&flat, &ids)?;
Ok(Reconstructed {
index,
snapshot_id,
git_sha: sha,
by_id,
hnsw: OnceLock::new(),
})
}
/// Reconstruct (memoized) + search in one call. Returns `(score, occurrence)`
/// best-first. The built index is cached so repeated queries against the same
/// snapshot don't rebuild from Parquet; new data resolves to a fresh
/// `snapshot_id` (a new `Uuid` per build) and rebuilds.
pub fn search(
wh: &IcebergWarehouse,
repo: &str,
git_sha: Option<&str>,
model_profile: &str,
query: &[f32],
k: usize,
) -> Result<Vec<(f32, Occurrence)>> {
let r = reconstruct_cached(wh, repo, git_sha, model_profile)?;
// Live scan-kernel selection (G2): `search_auto` uses the int8/VNNI
// fast-scan + exact-f32 rerank for large corpora (>= I8_SCAN_THRESHOLD rows)
// and the pure f32 scan below it — both return the exact cosine top-k, so
// this is results-identical to the old `search`, just faster on big indexes.
// `NORNIR_VECTOR_I8=0` forces the pure f32 path (operator escape hatch).
let hits = if i8_scan_enabled() {
r.index.search_auto(query, k)
} else {
r.index.search(query, k)
};
map_hits(&r, hits)
}
/// Map raw `(id, score)` kNN hits to `(score, Occurrence)` via the snapshot's
/// `by_id`. Shared by the exact [`search`] and the ANN [`dense_search_ann`].
fn map_hits(r: &Reconstructed, hits: Vec<(u64, f32)>) -> Result<Vec<(f32, Occurrence)>> {
let mut out = Vec::with_capacity(hits.len());
for (id, score) in hits {
let occ = r
.by_id
.get(&id)
.ok_or_else(|| anyhow!("hit id {id} has no occurrence"))?
.clone();
out.push((score, occ));
}
Ok(out)
}
/// **Dense kNN for the LIVE path** — routes through the HNSW ANN index when the
/// snapshot has one (large corpus, `NORNIR_VECTOR_HNSW` on), else the exact flat
/// scan. Used by [`hybrid_search`] so the shipped `vector_search` MCP tool gets
/// O(log n) retrieval on big codebases, with the flat scan as the exact oracle +
/// small-corpus fallback. `search` (the exact public API) is deliberately left on
/// the flat scan — it is the oracle every test/bench measures against.
fn dense_search_ann(
wh: &IcebergWarehouse,
repo: &str,
git_sha: Option<&str>,
model_profile: &str,
query: &[f32],
k: usize,
) -> Result<Vec<(f32, Occurrence)>> {
let r = reconstruct_cached(wh, repo, git_sha, model_profile)?;
let hits = match r.ann() {
Some(h) => h.search(query, k),
None if i8_scan_enabled() => r.index.search_auto(query, k),
None => r.index.search(query, k),
};
map_hits(&r, hits)
}
/// **Hybrid retrieval** — the LIVE Mimir SEARCH path (optimize cycle 1). Fuse the
/// dense cosine kNN of [`search`] with the persisted BM25 lexical index
/// ([`crate::index::Index`]) via Reciprocal-Rank-Fusion
/// ([`crate::knowledge::fusion::fuse_bm25_dense`]) — the productionized version of
/// the `nornir.mimir.search.hybrid_rrf_*` bench. Dense whiffs on rare exact
/// identifiers / error-strings; BM25 whiffs on paraphrase; fusing the two *ranked*
/// lists on RANK (no score calibration, no new dep) recovers the relevant files
/// EITHER engine buried, lifting recall at ~zero latency.
///
/// Both engines are pulled to a candidate pool wider than `k` so the fusion has
/// room to promote a file that only ONE engine ranked high. The two lists share
/// one key namespace — the repo-relative file path (`Occurrence::file` ==
/// `Hit::path`). The output is best-first by **fused RRF score**; a file only BM25
/// surfaced (no dense span) is emitted with a whole-file span so the recall gain
/// is actually delivered on the wire. A BM25 lookup failure (e.g. no text index
/// built yet) degrades cleanly to the dense-only ranking rather than erroring.
///
/// `query_text` is the raw query for BM25; `query_vec` is its embedding for the
/// dense arm. `k == 0` returns the full fused ranking.
pub fn hybrid_search(
wh: &IcebergWarehouse,
index: &crate::index::Index,
repo: &str,
git_sha: Option<&str>,
model_profile: &str,
query_text: &str,
query_vec: &[f32],
k: usize,
) -> Result<Vec<(f32, Occurrence)>> {
let pool = if k == 0 { 128 } else { (k * 8).max(64) };
// Dense arm — cosine kNN through the HNSW ANN index when present (O(log n) on
// large corpora), else the exact flat scan. Keep each file's BEST (first =
// top-rank) occurrence so the wire result can carry a real span, and build the
// dense rank list.
let dense_hits = dense_search_ann(wh, repo, git_sha, model_profile, query_vec, pool)?;
let mut occ_by_file: HashMap<String, Occurrence> = HashMap::new();
let mut dense_keys: Vec<String> = Vec::with_capacity(dense_hits.len());
for (_, occ) in &dense_hits {
if !occ_by_file.contains_key(&occ.file) {
dense_keys.push(occ.file.clone());
occ_by_file.insert(occ.file.clone(), occ.clone());
}
}
// Lexical arm — BM25 over the SAME repo, file-path keys. Degrade to dense-only
// on failure (no persisted text index yet) rather than failing the search.
let bm25_keys: Vec<String> = match index.search(
query_text,
Some(crate::index::Corpus::Code),
Some(repo),
pool,
) {
Ok(hits) => {
let mut v: Vec<String> = Vec::with_capacity(hits.len());
for h in hits {
if !v.contains(&h.path) {
v.push(h.path);
}
}
v
}
Err(_) => Vec::new(),
};
// Fuse on RANK — the bench-proven entrypoint. Equal weights (the tuned
// default); `k` truncates to the caller's top-k.
let fused = crate::knowledge::fusion::fuse_bm25_dense(
&bm25_keys,
&dense_keys,
1.0,
1.0,
crate::knowledge::fusion::DEFAULT_RRF_K,
k,
);
// Re-emit best-first with a monotone descending fused score. A BM25-only file
// (dense never surfaced it) gets a whole-file span — the file is the retrieval
// unit and dropping it would forfeit the recall gain the fusion exists for.
let mut out = Vec::with_capacity(fused.len());
for (rank, file) in fused.into_iter().enumerate() {
let score = 1.0 / (rank as f32 + 1.0);
let occ = occ_by_file.remove(&file).unwrap_or(Occurrence {
content_hash: String::new(),
file,
start_line: 0,
end_line: 0,
});
out.push((score, occ));
}
Ok(out)
}
/// **Batched** reconstruct + search: exact top-`k` for MANY query vectors against
/// the same snapshot, reading the reconstructed matrix **once** instead of once
/// per query (see [`znippy_zoomies::vann::VectorIndex::search_batch`]). `out[i]`
/// is the best-first `(score, occurrence)` list for `queries[i]`, byte-identical
/// to calling [`search`] per query — same exact cosine, same ranking. This is the
/// entry point for bulk work where you hold many queries at once: re-ranking a
/// candidate set, or a self-similarity / near-duplicate sweep. For a single query
/// use [`search`] (the batch amortization only pays off with many queries).
///
/// Every query vector must have the index's dimensionality.
pub fn search_batch(
wh: &IcebergWarehouse,
repo: &str,
git_sha: Option<&str>,
model_profile: &str,
queries: &[Vec<f32>],
k: usize,
) -> Result<Vec<Vec<(f32, Occurrence)>>> {
if queries.is_empty() {
return Ok(Vec::new());
}
let r = reconstruct_cached(wh, repo, git_sha, model_profile)?;
let dim = r.index.dim();
let mut flat = Vec::with_capacity(queries.len() * dim);
for (qi, q) in queries.iter().enumerate() {
ensure!(
q.len() == dim,
"query {qi} dim {} != index dim {dim}",
q.len()
);
flat.extend_from_slice(q);
}
// vann streams the corpus once, scoring each row against all queries — exact.
let per_q = r.index.search_batch(&flat, k);
let mut out = Vec::with_capacity(per_q.len());
for hits in per_q {
let mut one = Vec::with_capacity(hits.len());
for (id, score) in hits {
let occ = r
.by_id
.get(&id)
.ok_or_else(|| anyhow!("hit id {id} has no occurrence"))?
.clone();
one.push((score, occ));
}
out.push(one);
}
Ok(out)
}
/// Whether the int8/VNNI fast-scan path may be selected on the live search.
/// On by default (`VectorIndex::search_auto` still only uses it above its size
/// threshold); `NORNIR_VECTOR_I8=0`/`false`/`off` forces the pure f32 scan for
/// every corpus size — an escape hatch if a reviewer ever wants to A/B the two.
fn i8_scan_enabled() -> bool {
match std::env::var("NORNIR_VECTOR_I8") {
Ok(v) => !matches!(
v.trim().to_ascii_lowercase().as_str(),
"0" | "false" | "off" | "no"
),
Err(_) => true,
}
}
/// Whether the live dense-kNN path ([`dense_search_ann`], the `hybrid_search`
/// dense arm) may use the HNSW ANN index. On by default (still only built at/above
/// [`HNSW_MIN_ROWS`]); `NORNIR_VECTOR_HNSW=0`/`false`/`off`/`no` forces the exact
/// flat scan everywhere — the operator escape hatch to A/B ANN vs the oracle.
fn hnsw_enabled() -> bool {
match std::env::var("NORNIR_VECTOR_HNSW") {
Ok(v) => !matches!(
v.trim().to_ascii_lowercase().as_str(),
"0" | "false" | "off" | "no"
),
Err(_) => true,
}
}
// ----- index cache -----------------------------------------------------------
/// Most-recently-built indexes to keep resident. A **secondary** bound on lookup
/// cost only: the memory bound is [`INDEX_CACHE_BUDGET`]. This constant used to
/// be the ONLY bound, and it bounded nothing — one entry is a whole snapshot's
/// flat vectors plus, once a query has run, an HNSW graph that keeps its own
/// second copy of them.
const INDEX_CACHE_CAP: usize = 8;
/// The byte budget for [`INDEX_CACHE`] (`NORNIR_VECTOR_INDEX_MAX_BYTES`, default
/// 2 GiB) — the shared [`crate::cache_budget`] policy, the same one
/// [`crate::knowledge::query`] charges against.
pub static INDEX_CACHE_BUDGET: Budget = Budget::new(
"vector index",
"NORNIR_VECTOR_INDEX_MAX_BYTES",
2 * 1024 * 1024 * 1024,
);
/// Cache key: `(snapshot_id, model_profile)`.
#[derive(PartialEq, Eq, Clone)]
struct IndexKey {
snapshot_id: Uuid,
model_profile: String,
}
/// Process-wide cache of built indexes, keyed by `(snapshot_id, model_profile)`.
///
/// Correctness rests on snapshot immutability: every index build mints a fresh
/// `snapshot_id` (`Uuid::new_v4`) and a snapshot's manifest + content-addressed
/// vectors never change after it's written. So new data resolves to a *new*
/// `snapshot_id` (cache miss → rebuild) and a hit can never serve stale rows.
/// `model_profile` is part of the key so two profiles never share an entry
/// (a snapshot already belongs to exactly one profile; this makes that explicit).
///
/// Bounded in BYTES by [`INDEX_CACHE_BUDGET`], and **charged on use** — see
/// [`CachedIndex`], because an entry grows after admission.
static INDEX_CACHE: BudgetedCache<IndexKey, Reconstructed> =
BudgetedCache::new("vector index", &INDEX_CACHE_BUDGET, INDEX_CACHE_CAP);
/// An index served out of [`INDEX_CACHE`], which re-charges the cache with the
/// entry's TRUE size when the caller is done with it.
///
/// This type exists because a `Reconstructed` GROWS after it is admitted, inside
/// an `Arc` the cache has already handed out: the first ANN query builds the
/// HNSW graph in a `OnceLock` (and that graph keeps its own copy of the whole
/// `n × dim` float matrix), and the first int8 query builds the quantized matrix
/// in a second `OnceLock` inside `VectorIndex`. A budget charged only at insert
/// would under-count exactly the entries worth evicting — the hot ones. Making
/// the recharge a `Drop` puts it on ONE writer that every query path goes
/// through, rather than three call sites that must each remember.
struct CachedIndex {
key: IndexKey,
value: Arc<Reconstructed>,
}
impl std::ops::Deref for CachedIndex {
type Target = Reconstructed;
fn deref(&self) -> &Reconstructed {
&self.value
}
}
impl Drop for CachedIndex {
fn drop(&mut self) {
let bytes = self.value.resident_bytes();
INDEX_CACHE.recharge(|k| *k == self.key, bytes);
}
}
/// Resolve the snapshot for `(repo, git_sha, model_profile)` and return its
/// built index — from cache when present, else built once and memoized. The
/// snapshot resolution (a cheap scan of `embedding_snapshots`) runs every call,
/// so a query after new data always rebuilds.
fn reconstruct_cached(
wh: &IcebergWarehouse,
repo: &str,
git_sha: Option<&str>,
model_profile: &str,
) -> Result<CachedIndex> {
wh.block_on(async {
let (snapshot_id, sha) = resolve_snapshot(wh, repo, git_sha, model_profile).await?;
let key = IndexKey {
snapshot_id,
model_profile: model_profile.to_string(),
};
if let Some(hit) = INDEX_CACHE.get(|k| *k == key) {
return Ok(CachedIndex { key, value: hit });
}
let built = Arc::new(build_reconstructed(wh, snapshot_id, sha, model_profile).await?);
INDEX_CACHE.put(
key.clone(),
built.resident_bytes(),
built.clone(),
|k| *k == key,
);
Ok(CachedIndex { key, value: built })
})
}
/// Drop every built index. Test/bench hook — forces the cold rebuild path.
pub fn clear_index_cache() {
INDEX_CACHE.clear();
}
/// `(resident entries, resident bytes)` in the built-index cache.
pub fn index_cache_stats() -> (usize, usize) {
let s = INDEX_CACHE.stats();
(s.entries, s.bytes)
}
// ----- iceberg helpers -------------------------------------------------------
struct ManifestRow {
ordinal: i32,
content_hash: String,
file: String,
start_line: usize,
end_line: usize,
}
async fn find_snapshot(
wh: &IcebergWarehouse,
repo: &str,
git_sha: &str,
model_profile: &str,
chunker_hash: &str,
fileset_hash: &str,
) -> Result<Option<Uuid>> {
let table = wh
.catalog()
.load_table(&wh.table_ident(TABLE_EMBEDDING_SNAPSHOTS))
.await?;
let predicate = Reference::new("repo")
.equal_to(Datum::string(repo))
.and(Reference::new("git_sha").equal_to(Datum::string(git_sha)))
.and(Reference::new("model_profile").equal_to(Datum::string(model_profile)));
let scan = table
.scan()
.with_filter(predicate)
.select([
"snapshot_id",
"repo",
"git_sha",
"model_profile",
"chunker_hash",
"fileset_hash",
])
.build()?;
let batches: Vec<RecordBatch> = scan.to_arrow().await?.try_collect().await?;
for b in &batches {
let ids = col::<StringArray>(b, "snapshot_id")?;
let repos = col::<StringArray>(b, "repo")?;
let shas = col::<StringArray>(b, "git_sha")?;
let mps = col::<StringArray>(b, "model_profile")?;
let chs = col::<StringArray>(b, "chunker_hash")?;
let fhs = col::<StringArray>(b, "fileset_hash")?;
for i in 0..b.num_rows() {
if repos.value(i) == repo
&& shas.value(i) == git_sha
&& mps.value(i) == model_profile
&& chs.value(i) == chunker_hash
&& fhs.value(i) == fileset_hash
{
return Ok(Some(Uuid::parse_str(ids.value(i))?));
}
}
}
Ok(None)
}
pub(crate) async fn existing_content_hashes(
wh: &IcebergWarehouse,
model_profile: &str,
) -> Result<HashSet<String>> {
existing_content_hashes_in(wh, TABLE_EMBEDDINGS, model_profile).await
}
/// [`existing_content_hashes`] of another vector table of the same schema
/// (the dependency ring's: [`crate::warehouse::iceberg::TABLE_EMBEDDINGS_RING`]).
pub(crate) async fn existing_content_hashes_in(
wh: &IcebergWarehouse,
vector_table: &str,
model_profile: &str,
) -> Result<HashSet<String>> {
let table = wh
.catalog()
.load_table(&wh.table_ident(vector_table))
.await?;
let scan = table
.scan()
.with_filter(Reference::new("model_profile").equal_to(Datum::string(model_profile)))
.select(["content_hash", "model_profile"])
.build()?;
let batches: Vec<RecordBatch> = scan.to_arrow().await?.try_collect().await?;
let mut out = HashSet::new();
for b in &batches {
let hs = col::<StringArray>(b, "content_hash")?;
let mps = col::<StringArray>(b, "model_profile")?;
for i in 0..b.num_rows() {
if mps.value(i) == model_profile {
out.insert(hs.value(i).to_string());
}
}
}
Ok(out)
}
/// One writer of the vector tables at a time, process-wide. A consume's
/// `embeddings` feed and the background ring job ([`super::fnvec::ring`])
/// append to the same tables from two threads: every append takes this, so two
/// commits to one table never interleave. Held for one append, never for an embed.
pub(crate) static VECTOR_WRITE: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
pub(crate) async fn append_embeddings(
wh: &IcebergWarehouse,
model_profile: &str,
dim: usize,
normalized: bool,
rows: &[(&str, &[f32])],
) -> Result<()> {
append_embeddings_to(wh, TABLE_EMBEDDINGS, model_profile, dim, normalized, rows).await
}
/// [`append_embeddings`] into another vector table of the same schema.
pub(crate) async fn append_embeddings_to(
wh: &IcebergWarehouse,
vector_table: &str,
model_profile: &str,
dim: usize,
normalized: bool,
rows: &[(&str, &[f32])],
) -> Result<()> {
let _one_writer = VECTOR_WRITE.lock().await;
let table = wh
.catalog()
.load_table(&wh.table_ident(vector_table))
.await?;
let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
let elem = list_element_field(&schema, "vector")?;
let content: Vec<String> = rows.iter().map(|(h, _)| h.to_string()).collect();
let mps: Vec<String> = rows.iter().map(|_| model_profile.to_string()).collect();
let dims: Vec<i32> = rows.iter().map(|_| dim as i32).collect();
let norms: Vec<bool> = rows.iter().map(|_| normalized).collect();
let vhashes: Vec<String> = rows
.iter()
.map(|(_, v)| hex_sha256(f32_le_bytes(v).as_slice()))
.collect();
let mut vb = ListBuilder::new(Float32Builder::new()).with_field(elem);
for (_, v) in rows {
vb.values().append_slice(v);
vb.append(true);
}
let mut cols: Vec<Arc<dyn Array>> = vec![
Arc::new(StringArray::from(content)),
Arc::new(StringArray::from(mps)),
Arc::new(Int32Array::from(dims)),
Arc::new(vb.finish()),
Arc::new(BooleanArray::from(norms)),
Arc::new(StringArray::from(vhashes)),
];
// LANGUAGE AXIS + readable model id. Presence-guarded (the Arrow schema comes
// from the table's CURRENT schema and the column vector is positional), so a
// pre-axis `embeddings` table keeps its 6-column shape and a fresh one gets
// all 8. The walker only admits `.rs` files today, so `language` is `rust` by
// construction; `model_id` is the readable registry id behind the opaque
// `model_profile` sha256 — the key a cross-language cosine calibration must be
// stated against, since a threshold measured for one model says nothing about
// another. See `workspace_nordisk/.nornir/mimir-language-axis-design.md`.
let cur = table.metadata().current_schema();
if cur.field_by_name("language").is_some() {
let lang = vec![crate::warehouse::iceberg_schema::LANG_DEFAULT.to_string(); rows.len()];
cols.push(Arc::new(StringArray::from(lang)));
}
if cur.field_by_name("model_id").is_some() {
let ids = vec![crate::vector::embed_registry::selected_model_id(); rows.len()];
cols.push(Arc::new(StringArray::from(ids)));
}
let batch = RecordBatch::try_new(schema, cols)?;
append_batch(wh.catalog(), table, batch).await
}
async fn append_manifest(
wh: &IcebergWarehouse,
snapshot_id: Uuid,
model_profile: &str,
chunks: &[Chunk],
) -> Result<()> {
let table = wh
.catalog()
.load_table(&wh.table_ident(TABLE_EMBEDDING_MANIFEST))
.await?;
let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
let sid = snapshot_id.to_string();
let mut sids = Vec::with_capacity(chunks.len());
let mut ordinals = Vec::with_capacity(chunks.len());
let mut hashes = Vec::with_capacity(chunks.len());
let mut mps = Vec::with_capacity(chunks.len());
let mut files = Vec::with_capacity(chunks.len());
let mut starts = Vec::with_capacity(chunks.len());
let mut ends = Vec::with_capacity(chunks.len());
for (i, c) in chunks.iter().enumerate() {
sids.push(sid.clone());
ordinals.push(i as i32);
hashes.push(c.content_hash.clone());
mps.push(model_profile.to_string());
files.push(c.file.clone());
starts.push(c.start_line as i32);
ends.push(c.end_line as i32);
}
let cols: Vec<Arc<dyn Array>> = vec![
Arc::new(StringArray::from(sids)),
Arc::new(Int32Array::from(ordinals)),
Arc::new(StringArray::from(hashes)),
Arc::new(StringArray::from(mps)),
Arc::new(StringArray::from(files)),
Arc::new(Int32Array::from(starts)),
Arc::new(Int32Array::from(ends)),
];
let batch = RecordBatch::try_new(schema, cols)?;
append_batch(wh.catalog(), table, batch).await
}
async fn append_snapshot_row(
wh: &IcebergWarehouse,
snapshot_id: Uuid,
p: &IndexParams<'_>,
model_profile: &str,
fileset_hash: &str,
row_count: usize,
) -> Result<()> {
let table = wh
.catalog()
.load_table(&wh.table_ident(TABLE_EMBEDDING_SNAPSHOTS))
.await?;
let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
let cols: Vec<Arc<dyn Array>> = vec![
Arc::new(StringArray::from(vec![snapshot_id.to_string()])),
Arc::new(StringArray::from(vec![p.workspace.to_string()])),
Arc::new(StringArray::from(vec![p.repo.to_string()])),
Arc::new(StringArray::from(vec![p.git_sha.to_string()])),
Arc::new(StringArray::from(vec![p.branch.to_string()])),
Arc::new(StringArray::from(vec![model_profile.to_string()])),
Arc::new(StringArray::from(vec![p.chunker_hash.to_string()])),
Arc::new(StringArray::from(vec![fileset_hash.to_string()])),
Arc::new(BooleanArray::from(vec![p.complete])),
Arc::new(Int32Array::from(vec![row_count as i32])),
Arc::new(
TimestampMicrosecondArray::from(vec![Utc::now().timestamp_micros()])
.with_timezone("+00:00"),
),
];
let batch = RecordBatch::try_new(schema, cols)?;
append_batch(wh.catalog(), table, batch).await
}
async fn resolve_snapshot(
wh: &IcebergWarehouse,
repo: &str,
git_sha: Option<&str>,
model_profile: &str,
) -> Result<(Uuid, String)> {
let table = wh
.catalog()
.load_table(&wh.table_ident(TABLE_EMBEDDING_SNAPSHOTS))
.await?;
// Pushdown on repo+model only; match `git_sha` as a PREFIX in Rust so a
// short sha (e.g. the 12-char display form) resolves to the full snapshot.
let predicate = Reference::new("repo")
.equal_to(Datum::string(repo))
.and(Reference::new("model_profile").equal_to(Datum::string(model_profile)));
let scan = table
.scan()
.with_filter(predicate)
.select([
"snapshot_id",
"repo",
"git_sha",
"model_profile",
"ts_micros",
])
.build()?;
let batches: Vec<RecordBatch> = scan.to_arrow().await?.try_collect().await?;
let mut candidates: Vec<(Uuid, String, i64)> = Vec::new();
for b in &batches {
let ids = col::<StringArray>(b, "snapshot_id")?;
let repos = col::<StringArray>(b, "repo")?;
let shas = col::<StringArray>(b, "git_sha")?;
let mps = col::<StringArray>(b, "model_profile")?;
let ts = col::<TimestampMicrosecondArray>(b, "ts_micros")?;
for i in 0..b.num_rows() {
if repos.value(i) != repo || mps.value(i) != model_profile {
continue;
}
if let Some(sha) = git_sha {
if !shas.value(i).starts_with(sha) {
continue;
}
}
candidates.push((
Uuid::parse_str(ids.value(i))?,
shas.value(i).to_string(),
ts.value(i),
));
}
}
if candidates.is_empty() {
bail!(
"no embedding snapshot for repo `{repo}`{}",
git_sha.map(|s| format!(" at sha {s}")).unwrap_or_default()
);
}
candidates.sort_by_key(|c| c.2);
let chosen = candidates.pop().unwrap();
Ok((chosen.0, chosen.1))
}
async fn read_manifest(wh: &IcebergWarehouse, snapshot_id: Uuid) -> Result<Vec<ManifestRow>> {
let table = wh
.catalog()
.load_table(&wh.table_ident(TABLE_EMBEDDING_MANIFEST))
.await?;
let sid = snapshot_id.to_string();
let scan = table
.scan()
.with_filter(Reference::new("snapshot_id").equal_to(Datum::string(sid.clone())))
.select([
"snapshot_id",
"ordinal",
"content_hash",
"file",
"start_line",
"end_line",
])
.build()?;
let batches: Vec<RecordBatch> = scan.to_arrow().await?.try_collect().await?;
let mut out = Vec::new();
for b in &batches {
let sids = col::<StringArray>(b, "snapshot_id")?;
let ords = col::<Int32Array>(b, "ordinal")?;
let hs = col::<StringArray>(b, "content_hash")?;
let files = col::<StringArray>(b, "file")?;
let starts = col::<Int32Array>(b, "start_line")?;
let ends = col::<Int32Array>(b, "end_line")?;
for i in 0..b.num_rows() {
if sids.value(i) != sid {
continue;
}
out.push(ManifestRow {
ordinal: ords.value(i),
content_hash: hs.value(i).to_string(),
file: files.value(i).to_string(),
start_line: starts.value(i) as usize,
end_line: ends.value(i) as usize,
});
}
}
out.sort_by_key(|m| m.ordinal);
Ok(out)
}
/// Returns (content_hash → vector, dim). Only hashes in `needed` are kept.
pub(crate) async fn read_vectors(
wh: &IcebergWarehouse,
model_profile: &str,
needed: &HashSet<&str>,
) -> Result<(HashMap<String, Vec<f32>>, usize)> {
let (out, dim) = read_vectors_present(wh, model_profile, needed).await?;
if out.is_empty() {
bail!("no embeddings found for the snapshot's content hashes");
}
Ok((out, dim))
}
/// [`read_vectors`] for a reader to whom "none of these is stored" is an
/// answer, not an error (the per-function vectors: coverage is reported).
pub(crate) async fn read_vectors_present(
wh: &IcebergWarehouse,
model_profile: &str,
needed: &HashSet<&str>,
) -> Result<(HashMap<String, Vec<f32>>, usize)> {
read_vectors_present_in(wh, TABLE_EMBEDDINGS, model_profile, needed).await
}
/// [`read_vectors_present`] from another vector table of the same schema
/// (`embeddings_norm`, the clone signals' normalised bodies).
pub(crate) async fn read_vectors_present_in(
wh: &IcebergWarehouse,
vector_table: &str,
model_profile: &str,
needed: &HashSet<&str>,
) -> Result<(HashMap<String, Vec<f32>>, usize)> {
let table = wh
.catalog()
.load_table(&wh.table_ident(vector_table))
.await?;
let scan = table
.scan()
.with_filter(Reference::new("model_profile").equal_to(Datum::string(model_profile)))
.select(["content_hash", "model_profile", "dim", "vector"])
.build()?;
let batches: Vec<RecordBatch> = scan.to_arrow().await?.try_collect().await?;
let mut out: HashMap<String, Vec<f32>> = HashMap::with_capacity(needed.len());
let mut dim = 0usize;
for b in &batches {
let hs = col::<StringArray>(b, "content_hash")?;
let mps = col::<StringArray>(b, "model_profile")?;
let dims = col::<Int32Array>(b, "dim")?;
let vecs = col::<ListArray>(b, "vector")?;
for i in 0..b.num_rows() {
if mps.value(i) != model_profile {
continue;
}
let h = hs.value(i);
if !needed.contains(h) || out.contains_key(h) {
continue;
}
let row = vecs.value(i);
let fa = row
.as_any()
.downcast_ref::<Float32Array>()
.ok_or_else(|| anyhow!("vector element not Float32Array"))?;
let v: Vec<f32> = fa.values().to_vec();
dim = dims.value(i) as usize;
if v.len() != dim {
bail!("stored vector len {} != dim {dim} for {h}", v.len());
}
out.insert(h.to_string(), v);
}
}
Ok((out, dim))
}
// ----- small helpers ---------------------------------------------------------
fn ensure_dim(v: &[f32], dim: usize, chunk: &Chunk) -> Result<()> {
if v.len() != dim {
bail!(
"embedding for {}:{}-{} has dim {} != model dim {dim}",
chunk.file,
chunk.start_line,
chunk.end_line,
v.len()
);
}
Ok(())
}
/// Stable identity of the corpus+chunking: sorted (file, span, content_hash).
fn fileset_hash_of(chunks: &[Chunk]) -> String {
let mut keys: Vec<String> = chunks
.iter()
.map(|c| {
format!(
"{}:{}:{}:{}",
c.file, c.start_line, c.end_line, c.content_hash
)
})
.collect();
keys.sort();
let mut h = nornir_hash::Sha256Hasher::new();
for k in keys {
h.update(k.as_bytes());
h.update(b"\n");
}
h.finalize_hex()
}
fn f32_le_bytes(v: &[f32]) -> Vec<u8> {
let mut out = Vec::with_capacity(v.len() * 4);
for &f in v {
out.extend_from_slice(&f.to_le_bytes());
}
out
}
fn hex_sha256(bytes: &[u8]) -> String {
nornir_hash::sha256_hex(bytes)
}
/// Downcast a projected column by name (projection reorders/drops columns).
pub(crate) fn col<'a, T: 'static>(batch: &'a RecordBatch, name: &str) -> Result<&'a T> {
batch
.column_by_name(name)
.ok_or_else(|| anyhow!("projected batch missing column `{name}`"))?
.as_any()
.downcast_ref::<T>()
.ok_or_else(|| anyhow!("column `{name}` has unexpected arrow type"))
}
/// The inner element `Field` of a list column — carries the `PARQUET:field_id`
/// metadata that `RecordBatch::try_new` validates against the table schema.
fn list_element_field(schema: &ArrowSchema, name: &str) -> Result<FieldRef> {
let field = schema
.field_with_name(name)
.with_context(|| format!("schema missing list column `{name}`"))?;
match field.data_type() {
DataType::List(elem) | DataType::LargeList(elem) => Ok(elem.clone()),
other => bail!("column `{name}` expected List, got {other:?}"),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::vector::chunk::{ChunkOptions, hash_text};
use crate::warehouse::iceberg::IcebergWarehouse;
fn profile() -> ModelProfile {
ModelProfile {
model_name: "test-embed".into(),
weights_sha: "w0".into(),
tokenizer_sha: "t0".into(),
pooling: "mean".into(),
normalize: true,
dim: 8,
dtype: "f32".into(),
}
}
/// Build a `ModelProfile` from a registry entry the way `embed_support`
/// does — model name + dim from the registry, fixed-but-distinct
/// weight/tokenizer SHAs per model so two models never share a key. This is
/// the profile an embedder for that model would report; we don't need to
/// run a (multi-GB) ONNX model to assert the KEYING.
fn registry_profile(m: &crate::vector::embed_registry::EmbedModel) -> ModelProfile {
ModelProfile {
model_name: m.model_name.into(),
weights_sha: format!("weights-of-{}", m.id),
tokenizer_sha: format!("tok-of-{}", m.id),
pooling: "mean".into(),
normalize: true,
dim: m.dim,
dtype: "f32".into(),
}
}
/// LAW 1 — inject-and-assert. Selecting a model from the registry yields a
/// profile that reports *that model's* dim and a *distinct* `model_profile`
/// key; two different models therefore produce different warehouse keys, so
/// their vectors can't collide. We assert the KEY/profile logic on a temp
/// warehouse — no second ONNX model is run.
#[test]
fn registry_models_produce_distinct_profiles_and_keys() {
let _g = crate::cache_budget::exclusive();
use crate::vector::embed_registry as reg;
let jina = reg::by_id("jina-v2-base-code").expect("default registered");
let mini = reg::by_id("minilm-l6-v2").expect("alt registered");
// Dimension-aware: the profile reports the registry model's dim.
let pj = registry_profile(jina);
let pm = registry_profile(mini);
assert_eq!(pj.dim, 768, "jina is 768-dim");
assert_eq!(pm.dim, 384, "minilm is 384-dim");
assert_ne!(pj.dim, pm.dim, "different models, different dims");
// Distinct model_profile keys (the warehouse dedup/segregation key).
assert_ne!(
pj.id(),
pm.id(),
"different models MUST hash to different model_profile keys"
);
// And those distinct keys actually segregate rows in the warehouse:
// index under jina's profile, then a search under minilm's profile
// (its own dim) finds NOTHING — vectors never mix across models.
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let ch = chunker_hash(&ChunkOptions::default());
let chunks = vec![mk("a.rs", 1, "alpha")];
// Embed into the jina profile (dim 768).
index_snapshot(&wh, ¶ms("repo", "sha1", &pj, &ch), &chunks, |missing| {
Ok(missing.iter().map(|_| vec![0.0f32; pj.dim]).collect())
})
.unwrap();
// A search keyed by minilm's profile id has no snapshot → error, NOT a
// cross-model hit. (Query dim is minilm's 384.)
let q = vec![0.0f32; pm.dim];
let err = search(&wh, "repo", Some("sha1"), &pm.id(), &q, 1).unwrap_err();
assert!(
err.to_string().contains("no embedding snapshot"),
"minilm profile must not see jina's vectors: {err}"
);
// The original (jina) profile still resolves its own snapshot.
let qj = {
let mut v = vec![0.0f32; pj.dim];
v[0] = 1.0;
v
};
let hits = search(&wh, "repo", Some("sha1"), &pj.id(), &qj, 1).unwrap();
assert_eq!(hits.len(), 1, "jina profile still finds its own vector");
}
fn unit(axis: usize) -> Vec<f32> {
let mut v = vec![0.0f32; 8];
v[axis] = 1.0;
v
}
/// Deterministic test "embedding": map known text to a basis axis.
fn axis_of(text: &str) -> usize {
match text {
"alpha" | "same" => 0,
"bravo" => 1,
"charlie" => 2,
"bravo-2" => 3,
"delta" => 4,
_ => 7,
}
}
/// The `embed_missing` closure used by `index_snapshot` tests.
fn embed(missing: &[Chunk]) -> Result<Vec<Vec<f32>>> {
Ok(missing.iter().map(|c| unit(axis_of(&c.text))).collect())
}
fn mk(file: &str, line: usize, text: &str) -> Chunk {
Chunk {
file: file.into(),
start_line: line,
end_line: line,
content_hash: hash_text(text),
text: text.into(),
}
}
fn params<'a>(
repo: &'a str,
sha: &'a str,
model: &'a ModelProfile,
ch: &'a str,
) -> IndexParams<'a> {
IndexParams {
workspace: "ws",
repo,
git_sha: sha,
branch: "main",
model,
chunker_hash: ch,
complete: true,
}
}
struct TestEmbedder;
impl Embedder for TestEmbedder {
fn profile(&self) -> ModelProfile {
profile()
}
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
Ok(texts.iter().map(|t| unit(axis_of(t))).collect())
}
}
/// An embedder with a KNOWN cost per chunk, and a deliberately expensive FIRST
/// call — the shape a real backend has when it page-faults a freshly-mmapped
/// weights file. Lets the estimator's warm-up discard be asserted rather than
/// assumed.
struct PacedEmbedder {
first_call_extra: std::time::Duration,
per_chunk: std::time::Duration,
calls: std::sync::atomic::AtomicUsize,
}
impl Embedder for PacedEmbedder {
fn profile(&self) -> ModelProfile {
profile()
}
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let n = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if n == 0 {
std::thread::sleep(self.first_call_extra);
}
std::thread::sleep(self.per_chunk * texts.len() as u32);
Ok(texts.iter().map(|t| unit(axis_of(t))).collect())
}
}
fn paced(first_extra_ms: u64, per_chunk_ms: u64) -> PacedEmbedder {
PacedEmbedder {
first_call_extra: std::time::Duration::from_millis(first_extra_ms),
per_chunk: std::time::Duration::from_millis(per_chunk_ms),
calls: std::sync::atomic::AtomicUsize::new(0),
}
}
/// The whole point: a corpus far too big to try is priced from a sample small
/// enough to be free. 10 ms/chunk over 60_000 chunks is ~10 min, so the gate
/// sits right at the budget and the arithmetic is checkable by hand.
#[test]
fn estimate_extrapolates_a_sample_to_the_whole_corpus() {
let e = paced(0, 10);
let sample: Vec<&str> = vec!["a", "b", "c", "d", "e", "f"];
let eta = e.estimate(&sample, 60_000).unwrap();
// 5 chunks timed (6 minus the discarded warm-up), ~10 ms each ⇒ ~100 chunk/s.
assert_eq!(eta.sampled, 5);
assert_eq!(eta.total_chunks, 60_000);
assert!(
(30.0..=100.0).contains(&eta.rate),
"rate {} outside the band a 10 ms/chunk sleep can produce",
eta.rate
);
// Loose because sleep() only guarantees a LOWER bound, and CI is noisy.
let mins = eta.eta.as_secs_f64() / 60.0;
assert!(mins >= 10.0, "60k chunks at <=100/s cannot beat 10 min: {mins}");
}
/// The warm-up must not be billed to the per-chunk rate. A 2 s first call on a
/// 10 ms/chunk backend would halve the reported rate if it were counted — and
/// would penalise exactly the large-weights backends quantization is meant to
/// reward.
#[test]
fn estimate_does_not_bill_the_warmup_to_the_rate() {
let sample: Vec<&str> = vec!["a", "b", "c", "d", "e", "f"];
let cold = paced(2_000, 10).estimate(&sample, 1_000).unwrap();
let warm = paced(0, 10).estimate(&sample, 1_000).unwrap();
// A 2 s first-touch penalty is ~40x the entire timed sample; if it leaked
// into the measurement the two rates could not land within 2x of each other.
let ratio = warm.rate / cold.rate;
assert!(
(0.5..=2.0).contains(&ratio),
"warm-up leaked into the timed window: warm {} vs cold {} chunk/s",
warm.rate,
cold.rate
);
}
/// The gate is the reason the estimate exists: 10 minutes is the line.
#[test]
fn the_budget_gate_is_the_ten_minute_line() {
assert_eq!(EMBED_BUDGET.as_secs(), 600);
let e = paced(0, 10); // ~100 chunk/s
// 100 chunks ⇒ ~1 s. Comfortably inside.
assert!(e.estimate(&["a", "b", "c", "d"], 100).unwrap().within_budget());
// 6 million chunks at that rate is ~17 hours — the `dwarf` case.
let over = paced(0, 10).estimate(&["a", "b", "c", "d"], 6_000_000).unwrap();
assert!(!over.within_budget());
assert!(over.summary().contains("OVER"), "summary: {}", over.summary());
}
/// The gate refuses a job it has priced above the budget, and the refusal is
/// actionable: it names the measured rate, the projected time, and both ways out.
#[test]
fn the_gate_refuses_a_job_it_priced_over_budget() {
// 250 ms/chunk ⇒ ~4 chunk/s. 3000 chunks ⇒ ~750 s, over the 600 s budget.
let e = paced(0, 250);
let texts: Vec<&str> = vec!["some chunk of source text"; 3000];
let err = check_embed_budget(&e, &texts).expect_err("750 s must not pass a 600 s budget");
let msg = err.to_string();
assert!(msg.contains("over the"), "not a budget refusal: {msg}");
assert!(
msg.contains("potion-base-8m"),
"refusal must name the faster model: {msg}"
);
assert!(
msg.contains("NORNIR_EMBED_BUDGET=off"),
"refusal must name the escape hatch: {msg}"
);
}
/// A job that fits is priced and waved through, and the caller gets the number.
#[test]
fn the_gate_passes_and_reports_a_job_that_fits() {
let e = paced(0, 10); // ~100 chunk/s
let texts: Vec<&str> = vec!["chunk"; 1000]; // ~10 s
let eta = check_embed_budget(&e, &texts)
.expect("10 s is inside a 600 s budget")
.expect("a timeable backend must produce an Eta");
assert_eq!(eta.total_chunks, 1000);
assert!(eta.within_budget());
}
/// An instant backend is NOT blocked — this is the static-table path, the one that
/// must never be refused.
///
/// Asserts the OUTCOME (*not refused*), never the mechanism. `TestEmbedder` is
/// instant, so on a fast box the 7 timed chunks land under `estimate`'s 1 µs clock
/// floor and the gate declines to invent a number (`Ok(None)`); on a slower one the
/// same 7 chunks take a hair over it, `estimate` succeeds, and the gate waves the job
/// through carrying an `Eta` (`Ok(Some)`). **Both are passes**, and which one you get
/// is a property of the host clock, not of the gate.
///
/// MEASURED 2026-08-16, oden: asserting `.is_none()` failed 40/40 in isolated runs —
/// the sample took 1.112–1.233 µs, i.e. 1.11–1.23x the floor. A guard that goes red
/// for being 200 ns slow is reporting the weather, not a defect. What is worth
/// guarding is that an instant backend is never REFUSED, and that is what this now
/// says. Red-when-broken is intact: make the gate refuse an unpriceable job and
/// `is_ok()` fails.
#[test]
fn the_gate_does_not_block_a_backend_too_fast_to_time() {
let texts: Vec<&str> = vec!["chunk"; 5000];
let verdict = check_embed_budget(&TestEmbedder, &texts);
assert!(
verdict.is_ok(),
"an instant backend must pass, not be refused: {:?}",
verdict.as_ref().err()
);
// If the host clock DID resolve it, the number must describe this corpus and
// must fit — an instant backend that priced itself over budget would mean the
// extrapolation, not the timing, is wrong.
if let Ok(Some(eta)) = verdict {
assert_eq!(
eta.total_chunks, 5000,
"a priced job must be priced for the corpus it was handed"
);
assert!(
eta.within_budget(),
"an instant backend cannot be over a 600 s budget: {eta:?}"
);
}
}
/// `NORNIR_EMBED_BUDGET` parsing: the units an operator would actually type, the
/// spellings that switch the gate off, and a HARD ERROR on the ones they mistype.
///
/// Calls [`parse_embed_budget`] — the function the gate itself calls. Until
/// 2026-08-16 this test re-implemented the match in its own body and asserted its own
/// copy, on the (true) grounds that the env var is process-global while these tests
/// run in parallel. The cost of that workaround: deleting the real parser outright
/// left this test green. The parser is now pure and takes the string, so there is
/// nothing global to work around and nothing to re-implement.
#[test]
fn the_budget_env_parses_units_and_refuses_nonsense() {
// The units.
for (raw, want_secs) in [
("600", 600.0),
("600s", 600.0),
("10m", 600.0),
("2h", 7200.0),
(" 90S ", 90.0), // trimmed and case-folded, as an operator types it
("0.5h", 1800.0), // fractional, because the parse is f64 and says so
] {
let got = parse_embed_budget(raw)
.unwrap_or_else(|e| panic!("`{raw}` must parse: {e}"))
.unwrap_or_else(|| panic!("`{raw}` is a duration, not the off switch"));
assert_eq!(
got.as_secs_f64(),
want_secs,
"`{raw}` should be {want_secs}s"
);
}
// The off switch, every spelling the doc promises.
for raw in ["off", "none", "no", "0", "OFF", " None "] {
assert!(
parse_embed_budget(raw)
.unwrap_or_else(|e| panic!("`{raw}` must parse: {e}"))
.is_none(),
"`{raw}` must switch the gate OFF, not set a budget"
);
}
// The mistypes. Each must be an Err naming the variable — never a silent
// fallback to the default, which is the whole reason this is a hard error:
// an operator who typed `10min` meant to change the policy.
for raw in ["10min", "banana", "", " ", "-5", "1x", "inf", "NaN", "2hh"] {
let err = parse_embed_budget(raw)
.err()
.unwrap_or_else(|| panic!("`{raw}` is not a duration and must be REFUSED"));
assert!(
err.to_string().contains(NORNIR_EMBED_BUDGET_ENV),
"the refusal must name the variable so the operator can find it: {err}"
);
}
}
/// **The bitter end**: the gate is wired into the real batched writer, and a
/// refusal fails every member AND leaves the warehouse untouched. This is the
/// difference between a gate that exists and a gate that is enforced.
#[test]
fn index_repos_batched_refuses_over_budget_and_writes_nothing() {
let opts = ChunkOptions::default();
// 3000 distinct chunks at ~4 chunk/s ⇒ ~750 s, over the 600 s budget.
let n = 3000;
let files: Vec<(String, String)> = (0..n)
.map(|i| (format!("f{i}.rs"), format!("body{i}")))
.collect();
let input = RepoInput {
workspace: "ws".into(),
repo: "ra".into(),
git_sha: "sha1".into(),
branch: "main".into(),
complete: true,
files,
};
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let slow = paced(0, 250);
let res = index_repos_batched(&wh, &[input], &opts, &slow);
assert_eq!(res.len(), 1);
let err = res[0].as_ref().expect_err("an over-budget batch must not succeed");
assert!(
err.to_string().contains("over the"),
"member should carry the budget refusal, got: {err:#}"
);
// Nothing was embedded, so nothing was persisted. A refusal is not a
// partial run — that is the whole point of gating BEFORE the slice loop.
let mp = profile().id();
let stored = wh
.block_on(async { existing_content_hashes(&wh, &mp).await })
.unwrap();
assert!(
stored.is_empty(),
"a refused batch must write zero vectors, found {}",
stored.len()
);
}
/// **The 2026-08-26 cold-start claim, settled at the gate**: a backend whose
/// LOAD is expensive (big weights, first-touch page faults — modelled as an
/// expensive first `embed` call) but whose sustained per-chunk work is fast
/// must PASS the budget for a corpus its sustained rate fits. If the load
/// cost leaked into the extrapolated rate, this 2 s first call against a
/// 1 ms/chunk backend would price 60 000 chunks (~1 min of real work) as
/// hours and refuse — the exact defect class of the Qwen2 presets re-reading
/// 4.7 GB per message. Verified GREEN before the 2026-08-26 fixes: the
/// probe's warm-up discard was never the live culprit (that was the
/// single-threaded sample batch — see `gate_sample_size` and
/// `embed_support::fanout_rows_cap`); this pins it against regression.
#[test]
fn the_gate_passes_a_cold_start_backend_with_fast_chunks() {
let e = paced(2_000, 1); // 2 s load-shaped first call, 1 ms/chunk after
let texts: Vec<&str> = vec!["some chunk of source text"; 60_000];
let eta = check_embed_budget(&e, &texts)
.expect("60k chunks at ~1000 chunk/s (~1 min) must pass a 10 min budget")
.expect("a timeable backend must produce an Eta");
assert!(
eta.within_budget(),
"load cost leaked into the per-chunk rate: {eta:?}"
);
}
/// One chunk cannot be both the warm-up and the measurement.
#[test]
fn estimate_refuses_a_sample_that_is_all_warmup() {
let e = paced(0, 1);
assert!(e.estimate(&["only"], 1_000).is_err());
assert!(e.estimate(&[], 1_000).is_err());
}
/// A backend fast enough to finish under the clock's resolution must say so
/// rather than report an infinite rate and a zero-second ETA.
///
/// **This test used to race the host clock.** It handed `TestEmbedder` a
/// 2-chunk sample and `expect_err`ed the guard — which only fires if one
/// trivial embed lands under a microsecond. On a quiet box it does. MEASURED
/// 2026-09-10 on t14s under seven parallel cargo builds, the single timed
/// forward took 1.293 µs (`rate: 773 395 chunk/s`), the guard did not fire,
/// and this was the ONE red in 1188 lib tests — having found no defect at
/// all. The guard worked; it simply had nothing to refuse.
///
/// So the branch is now ENTERED, not hoped for: [`eta_from_sample`] takes the
/// elapsed seconds as a parameter, and the assertion is on the code rather
/// than on how busy the machine is.
#[test]
fn estimate_refuses_to_time_the_untimeable() {
for secs in [0.0, MIN_TIMEABLE_SECS, MIN_TIMEABLE_SECS / 2.0] {
let err = eta_from_sample(1, secs, 1_000)
.expect_err("a sample at or under the clock's resolution cannot be timed");
assert!(
err.to_string().contains("too fast to time"),
"unexpected at secs={secs}: {err}"
);
}
// A sample with nothing in it is a different refusal, and says so.
let err = eta_from_sample(0, 1.0, 1_000).expect_err("zero timed chunks is not a rate");
assert!(err.to_string().contains("all warm-up"), "unexpected: {err}");
// INVERSE — without it a green here is satisfied by a function that
// refuses everything. A timeable sample yields the arithmetic, exactly:
// 4 chunks in 2 s is 2 chunk/s, and 1000 chunks at that rate is 500 s.
let eta = eta_from_sample(4, 2.0, 1_000).expect("2 s for 4 chunks is timeable");
assert_eq!(eta.rate, 2.0, "{eta:?}");
assert_eq!(eta.sampled, 4, "{eta:?}");
assert_eq!(eta.total_chunks, 1_000, "{eta:?}");
assert_eq!(eta.eta, std::time::Duration::from_secs(500), "{eta:?}");
}
/// …and the guard is still REACHED through the real `estimate`, so splitting
/// the arithmetic out did not leave `estimate` timing without checking.
///
/// This asserts the WIRING, not the clock: whatever `TestEmbedder`'s forward
/// costs on this box, `estimate` must agree with `eta_from_sample` fed the
/// same numbers — a refusal, or an `Eta` whose rate is the one the arithmetic
/// says it is. There is no load under which that can flip.
#[test]
fn estimate_routes_its_verdict_through_the_one_guard() {
match TestEmbedder.estimate(&["a", "b"], 1_000) {
Err(e) => assert!(
e.to_string().contains("too fast to time"),
"the only refusal a working embedder can give here: {e}"
),
Ok(eta) => {
assert_eq!(eta.sampled, 1, "one chunk timed, one spent as warm-up: {eta:?}");
assert!(
eta.rate.is_finite() && eta.rate > 0.0,
"a returned Eta must carry a real rate, never an infinity: {eta:?}"
);
let round_trip = eta_from_sample(eta.sampled, 1.0 / eta.rate, eta.total_chunks)
.expect("a rate that was accepted must be re-derivable");
assert_eq!(
round_trip.total_chunks, eta.total_chunks,
"estimate and eta_from_sample are one function: {eta:?}"
);
}
}
}
/// An embedder that serves the first `ok_calls` batches and then hard-fails —
/// the crash-mid-populate oracle.
struct FlakyEmbedder {
ok_calls: std::sync::atomic::AtomicUsize,
}
impl Embedder for FlakyEmbedder {
fn profile(&self) -> ModelProfile {
profile()
}
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
use std::sync::atomic::Ordering;
if self
.ok_calls
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
.is_err()
{
return Err(anyhow!("simulated embedder death"));
}
Ok(texts.iter().map(|t| unit(axis_of(t))).collect())
}
}
/// RED-WHEN-REVERTED: the streamed, per-slice-durable embed.
///
/// Restore the old monolithic `embedder.embed(&all_texts)` + single
/// `append_embeddings` and BOTH halves of this test fail:
///
/// * **Slicing happens at all** — a corpus of `> EMBED_SLICE_CHUNKS` distinct
/// chunks must reach the embedder in more than one call. The monolithic
/// path makes exactly one.
/// * **Committed slices survive a crash** — when the embedder dies partway,
/// the vectors from the slices that already landed are STILL in
/// `embeddings`. Under the monolithic path a mid-run death persists
/// **zero** vectors, so the >55-minute t14s populate had to restart from
/// nothing. (The member results are still errors either way: no member's
/// snapshot is complete. Durability of the work done ≠ success.)
#[test]
fn batched_embed_streams_in_slices_and_keeps_committed_ones_after_a_crash() {
let opts = ChunkOptions::default();
// One distinct chunk per file, spilling past a single slice.
let n = EMBED_SLICE_CHUNKS + 100;
let files: Vec<(String, String)> = (0..n)
.map(|i| (format!("f{i}.rs"), format!("body{i}")))
.collect();
let input = |repo: &str| RepoInput {
workspace: "ws".into(),
repo: repo.into(),
git_sha: "sha1".into(),
branch: "main".into(),
complete: true,
files: files.clone(),
};
// (1) SLICED: a healthy run over > one slice still yields every vector.
let dir_ok = tempfile::tempdir().unwrap();
let wh_ok = IcebergWarehouse::open(dir_ok.path()).unwrap();
let res = index_repos_batched(&wh_ok, &[input("ra")], &opts, &TestEmbedder);
let snap = res[0].as_ref().unwrap();
assert_eq!(
snap.new_vectors, n,
"every distinct chunk embedded exactly once"
);
let mp = profile().id();
let stored_ok = wh_ok
.block_on(async { existing_content_hashes(&wh_ok, &mp).await })
.unwrap();
assert_eq!(
stored_ok.len(),
n,
"all {n} vectors persisted across slices"
);
// (2) CRASH after the FIRST slice: the embedder dies on the call that would
// have served slice #2.
//
// `+ GATE_PROBE_CALLS` because `check_embed_budget` runs first and spends
// real `embed` calls on its sample. Without it the embedder would die
// inside the GATE, the loop would never run, and this test would assert
// durability of a run that never started — passing or failing for reasons
// that have nothing to do with slicing.
let dir_bad = tempfile::tempdir().unwrap();
let wh_bad = IcebergWarehouse::open(dir_bad.path()).unwrap();
let flaky = FlakyEmbedder {
ok_calls: std::sync::atomic::AtomicUsize::new(1 + GATE_PROBE_CALLS),
};
let res = index_repos_batched(&wh_bad, &[input("ra")], &opts, &flaky);
let err = res[0].as_ref().unwrap_err().to_string();
assert!(
err.contains("simulated embedder death"),
"the member fails — a partial embed is NEVER a green snapshot (got {err})"
);
// THE POINT: the first slice's work is on disk. Monolithic ⇒ 0.
let stored_bad = wh_bad
.block_on(async { existing_content_hashes(&wh_bad, &mp).await })
.unwrap();
assert_eq!(
stored_bad.len(),
EMBED_SLICE_CHUNKS,
"exactly the first slice's vectors survived the crash (monolithic embed persists 0)"
);
assert!(
stored_bad.len() < n,
"and the run really did die partway ({}/{n})",
stored_bad.len()
);
// …so the RETRY is cheap: a fresh run re-embeds only what is missing.
let res2 = index_repos_batched(&wh_bad, &[input("ra")], &opts, &TestEmbedder);
let snap2 = res2[0].as_ref().unwrap();
assert_eq!(
snap2.new_vectors,
n - EMBED_SLICE_CHUNKS,
"resume embeds ONLY the chunks the crashed run never got to"
);
}
#[test]
fn index_then_search_roundtrips() {
let _g = crate::cache_budget::exclusive();
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let m = profile();
let ch = chunker_hash(&ChunkOptions::default());
let chunks = vec![
mk("a.rs", 1, "alpha"),
mk("b.rs", 1, "bravo"),
mk("c.rs", 1, "charlie"),
];
let snap = index_snapshot(&wh, ¶ms("repo", "sha1", &m, &ch), &chunks, embed).unwrap();
assert_eq!(snap.occurrences, 3);
assert_eq!(snap.new_vectors, 3);
// Query near axis 1 → b.rs wins.
let mut q = unit(1);
q[2] = 0.1;
let hits = search(&wh, "repo", Some("sha1"), &m.id(), &q, 1).unwrap();
assert_eq!(hits[0].1.file, "b.rs");
}
/// GAP 3, the row-side bound, MEASURED by batch count: the per-repo path
/// embeds `slice_rows` missing chunks per call and appends each slice before
/// the next call, so the vectors in flight never exceed one slice. With 10
/// missing chunks and a slice of 4 the embedder is called exactly 3 times,
/// the largest call holds 4 rows, every vector lands, and the budget gate
/// (`price`) sees the WHOLE missing set exactly once. A second snapshot over
/// the same content embeds nothing — everything from every slice persisted.
#[test]
fn per_repo_index_streams_missing_chunks_in_bounded_slices() {
let _g = crate::cache_budget::exclusive();
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let m = profile();
let ch = chunker_hash(&ChunkOptions::default());
let chunks: Vec<Chunk> = (0..10)
.map(|i| mk(&format!("f{i}.rs"), 1, &format!("text-{i}")))
.collect();
let calls = std::cell::Cell::new(0usize);
let max_rows = std::cell::Cell::new(0usize);
let priced = std::cell::Cell::new(0usize);
let snap = index_snapshot_streamed(
&wh,
¶ms("repo", "sha1", &m, &ch),
&chunks,
4,
|missing| {
priced.set(priced.get() + 1);
assert_eq!(missing.len(), 10, "priced over the WHOLE missing set");
Ok(())
},
|slice| {
calls.set(calls.get() + 1);
max_rows.set(max_rows.get().max(slice.len()));
embed(slice)
},
)
.unwrap();
assert_eq!(snap.occurrences, 10);
assert_eq!(snap.new_vectors, 10);
assert_eq!(calls.get(), 3, "ceil(10 / 4) embed calls");
assert_eq!(max_rows.get(), 4, "no call ever held more than one slice");
assert_eq!(priced.get(), 1, "the budget gate runs once, not per slice");
// Everything from every slice is durable: a new sha over the same content
// finds every hash stored and embeds nothing.
let again = index_snapshot(&wh, ¶ms("repo", "sha2", &m, &ch), &chunks, |missing| {
panic!("nothing should be missing, got {}", missing.len())
})
.unwrap();
assert_eq!(again.new_vectors, 0);
assert_eq!(warehouse_stats(&wh).unwrap().embeddings, 10);
}
/// The one-call door keeps its contract: `index_snapshot` (FnOnce) is the
/// streamed path with a `usize::MAX` slice, so the embedder sees every
/// missing chunk in ONE call — `unchanged_content_dedups_across_commits`
/// and the callers in `knowledge::context` / `mimir_slices` rely on that.
#[test]
fn index_snapshot_is_the_single_call_door() {
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let m = profile();
let ch = chunker_hash(&ChunkOptions::default());
let chunks: Vec<Chunk> = (0..10)
.map(|i| mk(&format!("f{i}.rs"), 1, &format!("text-{i}")))
.collect();
let calls = std::cell::Cell::new(0usize);
let snap = index_snapshot(&wh, ¶ms("repo", "sha1", &m, &ch), &chunks, |missing| {
calls.set(calls.get() + 1);
assert_eq!(missing.len(), 10);
embed(missing)
})
.unwrap();
assert_eq!(snap.new_vectors, 10);
assert_eq!(calls.get(), 1);
// The env override on the per-repo slice is read, positive-only.
assert_eq!(embed_slice_chunks(), EMBED_SLICE_CHUNKS);
}
#[test]
fn reindex_same_snapshot_is_idempotent() {
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let m = profile();
let ch = chunker_hash(&ChunkOptions::default());
let chunks = vec![mk("a.rs", 1, "alpha"), mk("b.rs", 1, "bravo")];
let s1 = index_snapshot(&wh, ¶ms("repo", "sha1", &m, &ch), &chunks, embed).unwrap();
let s2 = index_snapshot(&wh, ¶ms("repo", "sha1", &m, &ch), &chunks, embed).unwrap();
assert_eq!(
s1.snapshot_id, s2.snapshot_id,
"same identity → same snapshot"
);
assert_eq!(s2.new_vectors, 0, "nothing re-embedded");
}
#[test]
fn unchanged_content_dedups_across_commits() {
let _g = crate::cache_budget::exclusive();
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let m = profile();
let ch = chunker_hash(&ChunkOptions::default());
// sha1: two chunks.
let b1 = vec![mk("a.rs", 1, "alpha"), mk("b.rs", 1, "bravo")];
index_snapshot(&wh, ¶ms("repo", "sha1", &m, &ch), &b1, embed).unwrap();
// sha2: a.rs unchanged (same content), b.rs changed + a new file.
let b2 = vec![
mk("a.rs", 1, "alpha"), // identical content → no re-embed
mk("b.rs", 1, "bravo-2"), // changed
mk("d.rs", 1, "delta"), // new
];
// The closure must only be asked to embed the 2 missing chunks.
let s2 = index_snapshot(&wh, ¶ms("repo", "sha2", &m, &ch), &b2, |missing| {
assert_eq!(missing.len(), 2, "only changed+new chunks embedded");
embed(missing)
})
.unwrap();
assert_eq!(
s2.new_vectors, 2,
"only the changed + new chunk are embedded"
);
// Time-travel: searching sha1 still returns a.rs near axis 0.
let hits = search(&wh, "repo", Some("sha1"), &m.id(), &unit(0), 1).unwrap();
assert_eq!(hits[0].1.file, "a.rs");
assert_eq!(hits[0].1.content_hash, hash_text("alpha"));
}
#[test]
fn duplicate_content_in_two_files_returns_both() {
let _g = crate::cache_budget::exclusive();
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let m = profile();
let ch = chunker_hash(&ChunkOptions::default());
// Same text in two files → one stored vector, two occurrences.
let chunks = vec![mk("x.rs", 1, "same"), mk("y.rs", 9, "same")];
let snap = index_snapshot(&wh, ¶ms("repo", "sha1", &m, &ch), &chunks, embed).unwrap();
assert_eq!(snap.new_vectors, 1, "deduped to one vector");
assert_eq!(snap.occurrences, 2);
let hits = search(&wh, "repo", Some("sha1"), &m.id(), &unit(0), 2).unwrap();
let files: HashSet<&str> = hits.iter().map(|(_, o)| o.file.as_str()).collect();
assert!(
files.contains("x.rs") && files.contains("y.rs"),
"both occurrences"
);
}
#[test]
fn cache_reuses_index_and_new_snapshot_rebuilds() {
let _g = crate::cache_budget::exclusive();
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let m = profile();
let ch = chunker_hash(&ChunkOptions::default());
let b1 = vec![mk("a.rs", 1, "alpha")];
let s1 = index_snapshot(&wh, ¶ms("repo", "sha1", &m, &ch), &b1, embed).unwrap();
// First search builds + caches the index for sha1's snapshot; a second
// search against the same snapshot must serve the cached Arc.
let _ = search(&wh, "repo", Some("sha1"), &m.id(), &unit(0), 1).unwrap();
let key = IndexKey {
snapshot_id: s1.snapshot_id,
model_profile: m.id(),
};
let cached = INDEX_CACHE.get(|k| *k == key).expect("snapshot index cached");
let _ = search(&wh, "repo", Some("sha1"), &m.id(), &unit(0), 1).unwrap();
let cached2 = INDEX_CACHE.get(|k| *k == key).expect("still cached");
assert!(
Arc::ptr_eq(&cached, &cached2),
"same snapshot reuses the built index"
);
// New data → new git_sha → new snapshot_id → distinct cache key. The
// query at sha2 must see the new file (no stale sha1-only result).
let b2 = vec![mk("a.rs", 1, "alpha"), mk("d.rs", 1, "delta")];
let s2 = index_snapshot(&wh, ¶ms("repo", "sha2", &m, &ch), &b2, embed).unwrap();
assert_ne!(s1.snapshot_id, s2.snapshot_id);
let hits = search(&wh, "repo", Some("sha2"), &m.id(), &unit(4), 1).unwrap();
assert_eq!(hits[0].1.file, "d.rs", "rebuild picks up new data");
// Latest (git_sha = None) resolves to the newest snapshot, not the cached one.
let latest = search(&wh, "repo", None, &m.id(), &unit(4), 1).unwrap();
assert_eq!(latest[0].1.file, "d.rs");
}
/// LAW (inject-assert): the int8/VNNI fast-scan path that `store::search`
/// now selects on large corpora (`VectorIndex::search_auto`, above
/// `I8_SCAN_THRESHOLD`) must return results IDENTICAL to the exact f32
/// `VectorIndex::search` — same ids, same order, same exact scores. This
/// guards the wired live path against quantization ever silently changing a
/// search result. Built directly on `VectorIndex` (the engine the warehouse
/// index is) so it exercises the >4096-row crossover without a heavy
/// Iceberg build.
#[test]
fn wired_i8_search_matches_exact_f32() {
let _g = crate::cache_budget::exclusive();
use crate::vector::{I8_SCAN_THRESHOLD, VectorIndex};
let dim = 128;
let n = I8_SCAN_THRESHOLD + 129; // above the crossover → int8 rerank path
let mk =
|seed: f32| -> Vec<f32> { (0..dim).map(|i| ((i as f32 + 1.0) * seed).sin()).collect() };
let mut idx = VectorIndex::new(dim).unwrap();
let mut flat = Vec::with_capacity(n * dim);
let mut ids = Vec::with_capacity(n);
for r in 0..n as u64 {
flat.extend(mk(0.0007 + r as f32 * 0.00009));
ids.push(r * 3 + 5);
}
idx.add(&flat, &ids).unwrap();
assert!(idx.len() >= I8_SCAN_THRESHOLD, "must hit the int8 path");
for &qs in &[0.0007 + 900.0 * 0.00009, 0.42, 0.99] {
let q = mk(qs);
for k in [1usize, 10, 50] {
let exact = idx.search(&q, k);
let auto = idx.search_auto(&q, k); // what store::search dispatches to
assert_eq!(exact.len(), auto.len(), "k={k} qs={qs}");
for (a, b) in exact.iter().zip(&auto) {
assert_eq!(a.0, b.0, "id/rank differs k={k} qs={qs}");
assert_eq!(a.1.to_bits(), b.1.to_bits(), "score differs k={k} qs={qs}");
}
}
}
}
#[test]
fn index_repo_chunks_embeds_and_searches() {
let _g = crate::cache_budget::exclusive();
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let files = vec![
("a.rs".to_string(), "alpha".to_string()),
("b.rs".to_string(), "bravo".to_string()),
];
let snap = index_repo(
&wh,
&RepoRef {
workspace: "ws",
repo: "repo",
git_sha: "sha1",
branch: "main",
complete: true,
},
&files,
&ChunkOptions::default(),
&TestEmbedder,
)
.unwrap();
assert_eq!(snap.occurrences, 2);
assert_eq!(snap.new_vectors, 2);
let hits = search(&wh, "repo", Some("sha1"), &profile().id(), &unit(1), 1).unwrap();
assert_eq!(hits[0].1.file, "b.rs");
}
/// LAW (inject-assert): the flattened, GPU-batched multi-repo
/// [`index_repos_batched`] must produce a warehouse INDISTINGUISHABLE from
/// running [`index_repo`] once per member — same per-member occurrences, same
/// new-vector totals (deduped once across members), and the same search
/// results. This guards the populate embed parallelization against ever
/// diverging from the serial per-member path it replaced.
#[test]
fn index_repos_batched_matches_per_member_serial() {
let _g = crate::cache_budget::exclusive();
let opts = ChunkOptions::default();
// Two members; member A and member B SHARE the text "shared" (so global
// dedup across members must credit it exactly once), plus unique content.
let files_a = vec![
("a1.rs".to_string(), "alpha".to_string()),
("a2.rs".to_string(), "shared".to_string()),
];
let files_b = vec![
("b1.rs".to_string(), "bravo".to_string()),
("b2.rs".to_string(), "shared".to_string()),
];
// SERIAL: index_repo per member into wh_seq.
let dir_seq = tempfile::tempdir().unwrap();
let wh_seq = IcebergWarehouse::open(dir_seq.path()).unwrap();
let rr = |repo: &'static str| RepoRef {
workspace: "ws",
repo,
git_sha: "sha1",
branch: "main",
complete: true,
};
let sa = index_repo(&wh_seq, &rr("ra"), &files_a, &opts, &TestEmbedder).unwrap();
let sb = index_repo(&wh_seq, &rr("rb"), &files_b, &opts, &TestEmbedder).unwrap();
// Serial: ra embeds alpha+shared (2 new); rb embeds bravo only (shared
// already stored) → 1 new. Total 3 distinct vectors.
assert_eq!(sa.new_vectors, 2);
assert_eq!(sb.new_vectors, 1);
// BATCHED: both members in one call into wh_bat.
let dir_bat = tempfile::tempdir().unwrap();
let wh_bat = IcebergWarehouse::open(dir_bat.path()).unwrap();
let inputs = vec![
RepoInput {
workspace: "ws".into(),
repo: "ra".into(),
git_sha: "sha1".into(),
branch: "main".into(),
complete: true,
files: files_a.clone(),
},
RepoInput {
workspace: "ws".into(),
repo: "rb".into(),
git_sha: "sha1".into(),
branch: "main".into(),
complete: true,
files: files_b.clone(),
},
];
let res = index_repos_batched(&wh_bat, &inputs, &opts, &TestEmbedder);
let ba = res[0].as_ref().unwrap();
let bb = res[1].as_ref().unwrap();
// Same occurrences per member as the serial path.
assert_eq!(ba.occurrences, sa.occurrences);
assert_eq!(bb.occurrences, sb.occurrences);
// Same new-vector attribution: shared is credited to the FIRST member (ra),
// so ra=2 (alpha+shared), rb=1 (bravo) — total 3, deduped once, NOT 4.
assert_eq!(ba.new_vectors, 2, "ra: alpha + shared");
assert_eq!(
bb.new_vectors, 1,
"rb: bravo (shared already credited to ra)"
);
assert_eq!(
ba.new_vectors + bb.new_vectors,
3,
"no cross-member double count"
);
// Same search behavior: querying rb near the "bravo" axis returns b1.rs.
let prof = profile().id();
let seq_hit = search(
&wh_seq,
"rb",
Some("sha1"),
&prof,
&unit(axis_of("bravo")),
1,
)
.unwrap();
let bat_hit = search(
&wh_bat,
"rb",
Some("sha1"),
&prof,
&unit(axis_of("bravo")),
1,
)
.unwrap();
assert_eq!(seq_hit[0].1.file, "b1.rs");
assert_eq!(
bat_hit[0].1.file, seq_hit[0].1.file,
"batched search matches serial"
);
}
/// LAW (inject-assert): `store::search_batch` (the wired Lever-B entry point)
/// must return, for each query, EXACTLY what `store::search` returns for that
/// query alone — same occurrences, same exact score bits. Guards the batched
/// warehouse path against ever diverging from the single-query oracle.
#[test]
fn wired_batch_search_matches_per_query() {
let _g = crate::cache_budget::exclusive();
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let files = vec![
("a.rs".to_string(), "alpha".to_string()),
("b.rs".to_string(), "bravo".to_string()),
("c.rs".to_string(), "charlie".to_string()),
];
index_repo(
&wh,
&RepoRef {
workspace: "ws",
repo: "repo",
git_sha: "sha1",
branch: "main",
complete: true,
},
&files,
&ChunkOptions::default(),
&TestEmbedder,
)
.unwrap();
let prof = profile().id();
let q0 = unit(1);
let q1 = unit(2);
let batch = search_batch(
&wh,
"repo",
Some("sha1"),
&prof,
&[q0.clone(), q1.clone()],
2,
)
.unwrap();
assert_eq!(batch.len(), 2, "one result list per query");
let s0 = search(&wh, "repo", Some("sha1"), &prof, &q0, 2).unwrap();
let s1 = search(&wh, "repo", Some("sha1"), &prof, &q1, 2).unwrap();
for (b, s) in [(&batch[0], &s0), (&batch[1], &s1)] {
assert_eq!(b.len(), s.len(), "batch/solo len differ");
for (a, c) in b.iter().zip(s) {
assert_eq!(a.0.to_bits(), c.0.to_bits(), "score differs from per-query");
assert_eq!(a.1.file, c.1.file, "occurrence differs from per-query");
}
}
}
/// RED-WHEN-BROKEN — the LIVE search path. The shipped `vector_search` MCP
/// tool routes through [`hybrid_search`]; this plants a corpus where HALF the
/// relevant files are **semantic-only** (dense finds, BM25 misses) and HALF
/// are **lexical-only** (a rare token in an off-topic file — BM25 finds, dense
/// misses), then asserts the fused ranking's recall@k STRICTLY beats the
/// dense-only ranking the tool used to return. Unwire the fusion (revert the
/// handler to `search`, or break `fuse_bm25_dense`) and `hybrid > dense`
/// flips → this turns red. Same file-path key namespace as the
/// `nornir.mimir.search.hybrid_rrf_*` bench; runs under `cargo test --lib`.
#[test]
fn hybrid_search_beats_dense_only_on_planted_corpus() {
let _g = crate::cache_budget::exclusive();
use crate::index::{Corpus, Index};
use crate::knowledge::fusion::recall_at_k;
let repo = "hybrid-live";
let kw = "flux_capacitor_z9"; // the rare token ONLY the lexical docs carry
let tmp = tempfile::tempdir().unwrap();
let ws = tmp.path();
struct Doc {
file: &'static str,
axis: usize,
token: bool,
relevant: bool,
}
// axis 0 == the dense query topic. Noise leans toward axis 0 (below) so it
// outranks the cosine-0 lexical docs → dense-only reliably misses them.
let docs = [
Doc {
file: "sem_a.rs",
axis: 0,
token: false,
relevant: true,
},
Doc {
file: "sem_b.rs",
axis: 0,
token: false,
relevant: true,
},
Doc {
file: "lex_a.rs",
axis: 5,
token: true,
relevant: true,
},
Doc {
file: "lex_b.rs",
axis: 6,
token: true,
relevant: true,
},
Doc {
file: "noise_0.rs",
axis: 1,
token: false,
relevant: false,
},
Doc {
file: "noise_1.rs",
axis: 2,
token: false,
relevant: false,
},
Doc {
file: "noise_2.rs",
axis: 3,
token: false,
relevant: false,
},
Doc {
file: "noise_3.rs",
axis: 4,
token: false,
relevant: false,
},
];
let mut chunks: Vec<Chunk> = Vec::new();
let mut paths: Vec<std::path::PathBuf> = Vec::new();
let mut relevant: Vec<String> = Vec::new();
let mut axis_by_hash: std::collections::HashMap<String, usize> =
std::collections::HashMap::new();
for d in &docs {
let mut body = format!("fn work_{}() {{}}\n", d.file.replace('.', "_"));
if d.token {
body.push_str(&format!("// {kw} lives in an off-topic file\n"));
}
std::fs::write(ws.join(d.file), &body).unwrap();
let hash = hash_text(&format!("{}-{}", d.file, body.len()));
axis_by_hash.insert(hash.clone(), d.axis);
chunks.push(Chunk {
file: d.file.into(),
start_line: 1,
end_line: 2,
content_hash: hash,
text: body,
});
paths.push(ws.join(d.file));
if d.relevant {
relevant.push(d.file.into());
}
}
// Dense snapshot: embed each chunk on its planted axis; noise gets a small
// axis-0 lean so it outranks the (cosine-0) lexical docs in dense-only.
let m = profile();
let ch = chunker_hash(&ChunkOptions::default());
let wh = IcebergWarehouse::open(&ws.join(".wh")).unwrap();
index_snapshot(&wh, ¶ms(repo, "sha1", &m, &ch), &chunks, |missing| {
Ok(missing
.iter()
.map(|c| {
let axis = *axis_by_hash.get(&c.content_hash).unwrap_or(&7);
let mut v = vec![0.0f32; m.dim];
v[axis] = 1.0;
if (1..=4).contains(&axis) {
v[0] = 0.5; // noise leans toward the dense query axis
}
v
})
.collect())
})
.unwrap();
// BM25 index over the same files (Hit::path == Occurrence::file namespace).
let idx_dir = tempfile::tempdir().unwrap();
let index = Index::open_at(ws, idx_dir.path()).unwrap();
index.index_files(repo, Corpus::Code, &paths).unwrap();
let k = relevant.len(); // 4
let qv = unit(0); // dense query = topic axis 0
// Dense-only — what the tool USED to return.
let dense = search(&wh, repo, Some("sha1"), &m.id(), &qv, 64).unwrap();
let dense_keys: Vec<String> = dense.iter().map(|(_, o)| o.file.clone()).collect();
let r_dense = recall_at_k(&dense_keys, &relevant, k);
// Hybrid — what the tool returns NOW (BM25 query = the rare token).
let hybrid = hybrid_search(&wh, &index, repo, Some("sha1"), &m.id(), kw, &qv, k).unwrap();
let hybrid_keys: Vec<String> = hybrid.iter().map(|(_, o)| o.file.clone()).collect();
let r_hybrid = recall_at_k(&hybrid_keys, &relevant, k);
assert!(
r_hybrid > r_dense,
"LIVE hybrid recall@{k} ({r_hybrid}) must STRICTLY beat dense-only \
({r_dense}); hybrid={hybrid_keys:?} dense={dense_keys:?}"
);
assert!(
(r_hybrid - 1.0).abs() < 1e-6,
"hybrid recovers ALL relevant (recall 1.0), got {r_hybrid}: {hybrid_keys:?}"
);
assert!(
r_dense <= 0.5 + 1e-6,
"dense-only must leave the lexical-only docs behind (recall <= 0.5), got {r_dense}"
);
}
/// Well-spread deterministic pseudo-random vector (NOT the smooth near-tie
/// shape) so recall@k is a meaningful metric.
fn rand_vec(seed: u64, dim: usize) -> Vec<f32> {
let mut s = seed
.wrapping_mul(0x9E37_79B9_7F4A_7C15)
.wrapping_add(0xABCD_1234);
(0..dim)
.map(|_| {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
((s >> 11) as f64 / (1u64 << 53) as f64 * 2.0 - 1.0) as f32
})
.collect()
}
/// RED-WHEN-BROKEN — the **recall floor** for the LIVE ANN path. Plants a
/// corpus larger than [`HNSW_MIN_ROWS`] (so the HNSW ANN index engages),
/// confirms the dense arm actually routes through it, and asserts its
/// **recall@10 vs the exact flat-scan oracle** stays above the floor. Break
/// the HNSW graph (or the wiring) and recall collapses → this turns red.
/// Runs under `cargo test --lib`.
#[test]
fn live_dense_arm_hnsw_recall_floor_vs_flat_oracle() {
let _g = crate::cache_budget::exclusive();
let dim = 32usize;
let n = 5_000usize; // > HNSW_MIN_ROWS (4096) → the ANN path engages
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let m = ModelProfile {
model_name: "hnsw-recall-test".into(),
weights_sha: "w0".into(),
tokenizer_sha: "t0".into(),
pooling: "mean".into(),
normalize: true,
dim,
dtype: "f32".into(),
};
let ch = chunker_hash(&ChunkOptions::default());
let chunks: Vec<Chunk> = (0..n)
.map(|i| {
let text = format!("chunk-{i}");
Chunk {
file: format!("f{i}.rs"),
start_line: 1,
end_line: 1,
content_hash: hash_text(&text),
text,
}
})
.collect();
index_snapshot(&wh, ¶ms("repo", "sha1", &m, &ch), &chunks, |missing| {
Ok(missing
.iter()
.map(|c| {
let seed = c
.content_hash
.bytes()
.fold(0u64, |a, b| a.wrapping_mul(131).wrapping_add(b as u64));
rand_vec(seed, dim)
})
.collect())
})
.unwrap();
let mp = m.id();
// The ANN index must actually engage for this corpus size.
let r = reconstruct_cached(&wh, "repo", Some("sha1"), &mp).unwrap();
assert!(
r.ann().is_some(),
"corpus > HNSW_MIN_ROWS must build an HNSW"
);
// recall@10: fraction of the exact flat-scan top-10 the ANN path recovers.
let mut hit = 0usize;
let mut tot = 0usize;
for q in 0..60u64 {
let qv = rand_vec(9_000_000 + q, dim);
let oracle: HashSet<String> = search(&wh, "repo", Some("sha1"), &mp, &qv, 10)
.unwrap()
.into_iter()
.map(|(_, o)| o.file)
.collect();
let ann = dense_search_ann(&wh, "repo", Some("sha1"), &mp, &qv, 10).unwrap();
for (_, o) in ann {
if oracle.contains(&o.file) {
hit += 1;
}
}
tot += oracle.len();
}
let recall = hit as f64 / tot.max(1) as f64;
assert!(
recall >= 0.90,
"live HNSW dense-arm recall@10 {recall:.3} fell below the 0.90 floor vs the flat oracle"
);
}
// ── the byte budget on INDEX_CACHE ───────────────────────────────────────
/// Build a snapshot of `n` chunks for `repo` and return its model-profile id.
fn seed_snapshot(wh: &IcebergWarehouse, repo: &str, n: usize, dim: usize) -> String {
let m = ModelProfile {
model_name: "budget-test".into(),
weights_sha: "w0".into(),
tokenizer_sha: "t0".into(),
pooling: "mean".into(),
normalize: true,
dim,
dtype: "f32".into(),
};
let ch = chunker_hash(&ChunkOptions::default());
let chunks: Vec<Chunk> = (0..n)
.map(|i| {
let text = format!("{repo}-chunk-{i}");
Chunk {
file: format!("{repo}/f{i}.rs"),
start_line: 1,
end_line: 1,
content_hash: hash_text(&text),
text,
}
})
.collect();
index_snapshot(wh, ¶ms(repo, "sha1", &m, &ch), &chunks, |missing| {
Ok(missing
.iter()
.map(|c| {
let seed = c
.content_hash
.bytes()
.fold(0u64, |a, b| a.wrapping_mul(131).wrapping_add(b as u64));
rand_vec(seed, dim)
})
.collect())
})
.unwrap();
m.id()
}
/// **RED WITHOUT THE BUDGET.** `INDEX_CACHE_CAP` is 8, so a count cap admits
/// a tiny 64-row index and a 6000-row one side by side and calls that
/// "bounded". The fixture is deliberately lopsided by ~two orders of
/// magnitude in rows, because a guard built on same-size entries cannot tell
/// a byte budget from an entry count — the identity-value trap.
///
/// RED SEEN (2026-08-23) by dropping the byte term from
/// `BudgetedCache::enforce` so only the count cap remains: phase A fails with
/// 2 entries resident where 1 was required.
#[test]
fn index_cache_is_bounded_in_bytes_not_in_entries() {
let _g = crate::cache_budget::exclusive();
let dim = 32usize;
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
clear_index_cache();
INDEX_CACHE_BUDGET.set_for_test(0);
let mp_small = seed_snapshot(&wh, "small", 64, dim);
let mp_big = seed_snapshot(&wh, "big", 6_000, dim);
// Measure the two entries first — the assertion below is meaningless if
// they are the same size.
clear_index_cache();
drop(reconstruct_cached(&wh, "small", Some("sha1"), &mp_small).unwrap());
let small_bytes = index_cache_stats().1;
clear_index_cache();
drop(reconstruct_cached(&wh, "big", Some("sha1"), &mp_big).unwrap());
let big_bytes = index_cache_stats().1;
assert!(
big_bytes > small_bytes * 10,
"the fixture must be lopsided: big {big_bytes} B vs small {small_bytes} B — \
a same-size fixture proves nothing about a BYTE budget"
);
// PHASE A — EVICTION. Budget with room for the big entry plus a little,
// but not for both. The big one must PUSH THE SMALL ONE OUT.
//
// (This phase exists because the first version of this guard was BLIND:
// it pinned the budget below the big entry, so the big one was refused
// at admission and the test passed even with eviction disabled entirely.
// Refusal and eviction are two different paths and a guard that only
// walks one of them proves only one of them.)
clear_index_cache();
let budget_a = big_bytes + small_bytes / 2;
INDEX_CACHE_BUDGET.set_for_test(budget_a);
drop(reconstruct_cached(&wh, "small", Some("sha1"), &mp_small).unwrap());
assert_eq!(index_cache_stats().0, 1, "the small index must be resident");
drop(reconstruct_cached(&wh, "big", Some("sha1"), &mp_big).unwrap());
let (entries, bytes) = index_cache_stats();
assert!(
bytes <= budget_a,
"{bytes} B resident in {entries} entries must fit the {budget_a} B budget \
(a count cap of {INDEX_CACHE_CAP} would have held BOTH)"
);
assert_eq!(
entries, 1,
"the big index must have EVICTED the small one, not sat beside it"
);
// And the evicted one is a genuine miss: served again by REBUILDING, not
// from a ghost.
let again = reconstruct_cached(&wh, "small", Some("sha1"), &mp_small).unwrap();
assert_eq!(again.index.len(), 64, "the re-decode must be the real index");
drop(again);
// PHASE B — REFUSAL. A budget below ONE entry: the big index is served
// to its caller and never retained, so it cannot pin the cache over
// budget forever.
clear_index_cache();
INDEX_CACHE_BUDGET.set_for_test(small_bytes * 2);
let refused_before = INDEX_CACHE_BUDGET.refused_too_big();
let served = reconstruct_cached(&wh, "big", Some("sha1"), &mp_big).unwrap();
assert_eq!(
served.index.len(),
6_000,
"an unaffordable index must still be SERVED, in full"
);
drop(served);
assert_eq!(
index_cache_stats().0,
0,
"an entry larger than the whole budget must not be retained"
);
assert_eq!(
INDEX_CACHE_BUDGET.refused_too_big(),
refused_before + 1,
"the refusal must be COUNTED — a cap-hit is loud, never silent"
);
INDEX_CACHE_BUDGET.set_for_test(0);
clear_index_cache();
}
/// **CHARGE ON USE.** An entry GROWS after admission: the first ANN query
/// builds the HNSW graph inside the already-shared `Arc`, and that graph
/// keeps its own copy of the whole float matrix. A budget charged only at
/// insert never sees it.
///
/// RED SEEN (2026-08-23) by deleting `impl Drop for CachedIndex`: the
/// charged figure stays at the admission value and this assertion fails with
/// after == before.
#[test]
fn a_query_that_builds_the_hnsw_is_charged_to_the_budget() {
let _g = crate::cache_budget::exclusive();
let dim = 32usize;
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
clear_index_cache();
// Generous budget: this test is about ACCOUNTING, not eviction.
INDEX_CACHE_BUDGET.set_for_test(4 * 1024 * 1024 * 1024);
// > HNSW_MIN_ROWS so the ANN path actually engages.
let mp = seed_snapshot(&wh, "grow", 5_000, dim);
let key = {
let held = reconstruct_cached(&wh, "grow", Some("sha1"), &mp).unwrap();
held.key.clone()
};
// Ask about THIS entry, not the whole cache: a total is at the mercy of
// anything else sharing the process.
let before = INDEX_CACHE.charge_of(|k| *k == key).expect("admitted");
assert!(before > 0, "an admitted index must be charged > 0 bytes");
// One ANN query — this is what builds the graph.
let q = rand_vec(4242, dim);
let hits = dense_search_ann(&wh, "grow", Some("sha1"), &mp, &q, 10).unwrap();
assert!(!hits.is_empty(), "the query must actually have run");
let after = INDEX_CACHE
.charge_of(|k| *k == key)
.expect("the entry itself must still be there");
assert!(
after > before,
"the HNSW graph built by the query ({before} B → {after} B) was never \
charged — a budget that only weighs at insert is blind to the larger \
half of a hot entry"
);
INDEX_CACHE_BUDGET.set_for_test(0);
clear_index_cache();
}
}