Skip to main content

ratel_ai_core/
dense_cache.rs

1//! Shared incremental dense-embedding cache backing the registries' semantic and
2//! hybrid engines.
3//!
4//! [`crate::ToolRegistry`] and [`crate::SkillRegistry`] rank two different item
5//! types (tools, skills) but embed and cosine-rank them identically: an **id-keyed
6//! map** of per-item vectors, extended incrementally. This type owns that
7//! structure and its operations so the two registries share one implementation
8//! instead of two copies that could drift. The registries keep their own
9//! trace-emitting search wrappers, since the trace event shapes differ (tool vs
10//! skill). See ADR-0011.
11//!
12//! Keying by id (not by position) is what lets `register` replace an item in
13//! place: on replace the registry calls [`DenseCache::invalidate`] to drop the
14//! stale vector, and the next [`DenseCache::extend`] re-embeds that id like any
15//! other missing one — so a re-registered id never leaves a stale embedding
16//! behind (RAT-378). The same keying lets `replace_all` drop the vector of an
17//! id that leaves the corpus outright, which no `extend` then re-embeds.
18
19use std::collections::HashMap;
20use std::sync::{Arc, Mutex, RwLock};
21
22use crate::dense_search::dense_search;
23use crate::embedding::{Embedder, EmbedderError, embedder_with_telemetry};
24use crate::embedding_artifact::{
25    ArtifactEntry, ArtifactEntryKind, ArtifactError, build_empty_artifact, hash_projection_text,
26    load_and_validate,
27};
28use crate::embedding_config::EmbeddingModel;
29use crate::trace::{TraceEvent, TraceSink};
30
31/// Single-input placeholder for the Endpoint identity probe.
32const ARTIFACT_WARM_PROBE_TEXT: &str = "__ratel_artifact_probe__";
33
34/// Result of [`DenseCache::warm_from_artifact`]: which corpus ids were taken from
35/// the artifact and which still need embedding (or another policy) by the caller.
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub(crate) struct WarmOutcome {
38    /// Corpus ids whose artifact vectors were committed.
39    pub reused: Vec<String>,
40    /// Corpus ids with no usable artifact entry (absent or projection_hash mismatch).
41    pub missing: Vec<String>,
42}
43
44/// Failure loading or applying an embedding artifact into the dense cache.
45#[derive(Debug, Clone)]
46pub enum WarmError {
47    /// Binary artifact failed [`crate::embedding_artifact::load_and_validate`].
48    Artifact(ArtifactError),
49    /// RAT1 header fingerprint does not match the configured embedder.
50    ArtifactModelMismatch {
51        /// Fingerprint recorded in the RAT1 header.
52        artifact: String,
53        /// Artifact-compat identity of the configured embedder.
54        active: String,
55    },
56    /// Embedder load, identity probe, or dimension check failed during warm.
57    Embedder(EmbedderError),
58}
59
60impl std::fmt::Display for WarmError {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        match self {
63            WarmError::Artifact(e) => write!(f, "{e}"),
64            WarmError::ArtifactModelMismatch { artifact, active } => write!(
65                f,
66                "embedding artifact was built with model {artifact}, but the configured embedding model is {active} — the artifact cannot warm this catalog (hint: rebuild the artifact with the current model, or configure the model the artifact was built with; artifact identities are opaque build-time values, so compare them rather than reading them)"
67            ),
68            WarmError::Embedder(e) => write!(f, "{e}"),
69        }
70    }
71}
72
73impl std::error::Error for WarmError {}
74
75impl From<ArtifactError> for WarmError {
76    fn from(value: ArtifactError) -> Self {
77        Self::Artifact(value)
78    }
79}
80
81impl From<EmbedderError> for WarmError {
82    fn from(value: EmbedderError) -> Self {
83        Self::Embedder(value)
84    }
85}
86
87/// An item the cache can embed and rank: its stable id and the flat searchable
88/// text fed to the embedder. Implemented by `Tool` and `Skill` in their
89/// respective registries, so the cache stays agnostic to the item type.
90pub(crate) trait Embeddable {
91    fn embed_id(&self) -> &str;
92    fn embed_text(&self) -> String;
93}
94
95/// A dense ranking paired with the query vector that produced it — see
96/// [`DenseCache::search_returning_query_vec`], which lets the usage-ranking arm
97/// reuse the embedding instead of paying for a second inference.
98pub(crate) type RankedWithQuery = (Vec<(String, f32)>, Vec<f32>);
99
100/// Per-item dense vectors, keyed by item id.
101#[derive(Default)]
102struct DenseCacheState {
103    /// `vectors[id]` is the embedding for the registry item with that id.
104    vectors: HashMap<String, Vec<f32>>,
105    /// Resolved identity of the model that produced every vector in `vectors`.
106    built_fingerprint: Option<String>,
107    /// Shared width of every vector in `vectors`.
108    dim: Option<usize>,
109}
110
111/// Per-item dense vectors, keyed by item id.
112pub(crate) struct DenseCache {
113    /// Vectors, dimension, and vector identity move together under one mutex so a
114    /// reader can never observe a partially-committed embedding batch.
115    state: Mutex<DenseCacheState>,
116    /// Builds take the write side; searches take the read side across query
117    /// embedding and ranking. This keeps one search in one vector space while
118    /// still allowing independent searches to run concurrently.
119    operation_lock: RwLock<()>,
120    /// Which embedding model backs this cache. Chosen per catalog; drives which
121    /// embedder [`Self::resolve_embedder`] loads. `Default` = built-in bge-small.
122    model: EmbeddingModel,
123    /// Test-only embedder override (`None` → the `model`'s embedder, loaded
124    /// lazily on first use). Lets tests inject a deterministic/failing embedder
125    /// without touching the network.
126    embedder_override: Option<Arc<dyn Embedder>>,
127}
128
129impl DenseCache {
130    pub(crate) fn new() -> Self {
131        Self::with_model(EmbeddingModel::Default)
132    }
133
134    /// A cache backed by an explicit embedding model (the configurable-model
135    /// path). The model is resolved lazily on first embed.
136    pub(crate) fn with_model(model: EmbeddingModel) -> Self {
137        Self {
138            state: Mutex::new(DenseCacheState::default()),
139            operation_lock: RwLock::new(()),
140            model,
141            embedder_override: None,
142        }
143    }
144
145    #[cfg(test)]
146    pub(crate) fn with_embedder(embedder: Arc<dyn Embedder>) -> Self {
147        Self::with_embedder_and_model(embedder, EmbeddingModel::Default)
148    }
149
150    /// Test helper that tags the cache with an [`EmbeddingModel`] variant while
151    /// still injecting a stub embedder — needed to exercise Endpoint warm probe
152    /// logic without a real network client.
153    #[cfg(test)]
154    pub(crate) fn with_embedder_and_model(
155        embedder: Arc<dyn Embedder>,
156        model: EmbeddingModel,
157    ) -> Self {
158        Self {
159            state: Mutex::new(DenseCacheState::default()),
160            operation_lock: RwLock::new(()),
161            model,
162            embedder_override: Some(embedder),
163        }
164    }
165
166    /// The embedder to use: an injected one (tests) or the configured model's,
167    /// whose one-time load telemetry is recorded on `sink`.
168    fn resolve_embedder(&self, sink: &dyn TraceSink) -> Result<Arc<dyn Embedder>, EmbedderError> {
169        match &self.embedder_override {
170            Some(e) => Ok(e.clone()),
171            None => {
172                self.model.validate()?;
173                embedder_with_telemetry(&self.model, sink)
174            }
175        }
176    }
177
178    /// Error unless the cache covers the whole corpus (`corpus_len` distinct ids).
179    /// A semantic/hybrid search never embeds inside the search path — the caller
180    /// must have built first, so no search silently pays the embedding cost.
181    /// Loads no model.
182    ///
183    /// The registries key their corpus by id, so `corpus_len` is the number of
184    /// distinct ids and `vectors.len()` is the number of embedded ids; the two
185    /// match exactly once every id is embedded. A re-register that
186    /// [`Self::invalidate`]s an id drops `vectors.len()` below `corpus_len`, so a
187    /// search after churn correctly reports `EmbeddingsNotBuilt` until the next
188    /// [`Self::extend`] — the same "build after registering" contract as a fresh
189    /// register.
190    /// The output width of the cached vectors, or `None` before any embedding
191    /// has run. Equals the active model's output dimension.
192    pub(crate) fn dim(&self) -> Option<usize> {
193        self.state.lock().expect("dense cache mutex poisoned").dim
194    }
195
196    /// The fingerprint of the model that built the current cache, or `None`
197    /// before any embedding has run. After a successful search this is the active
198    /// model's identity (query and corpus must agree, or `embed_query` errored),
199    /// so it doubles as "the model the query was embedded with" for the intent
200    /// graph's model check.
201    pub(crate) fn built_fingerprint(&self) -> Option<String> {
202        self.state
203            .lock()
204            .expect("dense cache mutex poisoned")
205            .built_fingerprint
206            .clone()
207    }
208
209    /// Embed arbitrary texts under the active model, returning the vectors and
210    /// the model fingerprint. Used to re-embed an intent graph's members when
211    /// rebuilding its centroids after a model change.
212    pub(crate) fn embed_texts_with_identity(
213        &self,
214        texts: &[String],
215        sink: &dyn TraceSink,
216    ) -> Result<(Vec<Vec<f32>>, String), EmbedderError> {
217        if texts.is_empty() {
218            return Ok((Vec::new(), self.built_fingerprint().unwrap_or_default()));
219        }
220        let embedder = self.resolve_embedder(sink)?;
221        let embedded = embedder.embed_batch_with_identity(texts)?;
222        Ok((embedded.value, embedded.fingerprint))
223    }
224
225    pub(crate) fn require_built(&self, corpus_len: usize) -> Result<(), EmbedderError> {
226        let cached = self
227            .state
228            .lock()
229            .expect("dense cache mutex poisoned")
230            .vectors
231            .len();
232        if cached < corpus_len {
233            return Err(EmbedderError::EmbeddingsNotBuilt);
234        }
235        Ok(())
236    }
237
238    /// Run `f` while holding the dense operation write lock. Used by registries to
239    /// compose warm + optional extend without releasing between steps (and without
240    /// re-entering [`Self::extend`], which would deadlock on this non-reentrant lock).
241    pub(crate) fn with_operation_write<R>(&self, f: impl FnOnce(&Self) -> R) -> R {
242        let _guard = self
243            .operation_lock
244            .write()
245            .expect("dense operation lock poisoned");
246        f(self)
247    }
248
249    /// Load vectors from a build-time embedding artifact for corpus ids whose
250    /// projection text still matches, without running inference on those ids.
251    ///
252    /// Matching (id + projection hash) runs before any embedder load. If nothing
253    /// can be reused, returns immediately with all corpus ids in `missing`.
254    /// Model identity is checked only when there is at least one reuse candidate
255    /// (Local/HF via [`Embedder::artifact_identity`]; Endpoint may probe once).
256    /// `built_fingerprint` continues to store the **runtime** identity.
257    /// Does not call [`Self::extend`] — the caller decides how to cover `missing`.
258    pub(crate) fn warm_from_artifact<'a, T: Embeddable + 'a>(
259        &self,
260        bytes: &[u8],
261        expected_kind: ArtifactEntryKind,
262        items: impl IntoIterator<Item = &'a T>,
263        sink: &dyn TraceSink,
264    ) -> Result<WarmOutcome, WarmError> {
265        self.with_operation_write(|cache| {
266            cache.warm_from_artifact_locked(bytes, expected_kind, items, sink)
267        })
268    }
269
270    /// Like [`Self::warm_from_artifact`], but assumes the caller already holds
271    /// [`Self::operation_lock`] for write.
272    pub(crate) fn warm_from_artifact_locked<'a, T: Embeddable + 'a>(
273        &self,
274        bytes: &[u8],
275        expected_kind: ArtifactEntryKind,
276        items: impl IntoIterator<Item = &'a T>,
277        sink: &dyn TraceSink,
278    ) -> Result<WarmOutcome, WarmError> {
279        let (header, entries) = load_and_validate(bytes)?;
280        // Known other kinds are ignored so one mixed RAT1 can warm either registry.
281        let by_id: HashMap<&str, &ArtifactEntry> = entries
282            .iter()
283            .filter(|e| e.kind == expected_kind)
284            .map(|e| (e.id.as_str(), e))
285            .collect();
286
287        let mut reused: Vec<(String, Vec<f32>)> = Vec::new();
288        let mut missing: Vec<String> = Vec::new();
289        for item in items {
290            let id = item.embed_id();
291            let text = item.embed_text();
292            match by_id.get(id) {
293                Some(entry) if entry.projection_hash == hash_projection_text(&text) => {
294                    reused.push((id.to_string(), entry.vector.clone()));
295                }
296                _ => missing.push(id.to_string()),
297            }
298        }
299
300        if reused.is_empty() {
301            return Ok(WarmOutcome {
302                reused: Vec::new(),
303                missing,
304            });
305        }
306
307        let embedder = self.resolve_embedder(sink)?;
308        let mut active_artifact = embedder.artifact_identity()?;
309        let mut runtime_identity = embedder.fingerprint();
310        if matches!(self.model, EmbeddingModel::Endpoint { .. })
311            && active_artifact != header.model_fingerprint
312        {
313            let probed = embedder.embed_batch_with_identity(&[ARTIFACT_WARM_PROBE_TEXT.into()])?;
314            active_artifact = probed.fingerprint.clone();
315            runtime_identity = probed.fingerprint;
316        }
317        if active_artifact != header.model_fingerprint {
318            let artifact = header.model_fingerprint.clone();
319            sink.record(TraceEvent::EmbedderModelMismatch {
320                built: artifact.clone(),
321                active: active_artifact.clone(),
322            });
323            return Err(WarmError::ArtifactModelMismatch {
324                artifact,
325                active: active_artifact,
326            });
327        }
328
329        let staged: Vec<Vec<f32>> = reused.iter().map(|(_, v)| v.clone()).collect();
330        let existing_dim = self.state.lock().expect("dense cache mutex poisoned").dim;
331        let expected_dim = validate_batch(
332            &staged,
333            reused.len(),
334            Some(existing_dim.unwrap_or(header.dim)),
335            &header.model_fingerprint,
336        )?;
337
338        let mut state = self.state.lock().expect("dense cache mutex poisoned");
339        if let Some(built) = &state.built_fingerprint
340            && built != &runtime_identity
341        {
342            let built = built.clone();
343            let active = runtime_identity.clone();
344            sink.record(TraceEvent::EmbedderModelMismatch {
345                built: built.clone(),
346                active: active.clone(),
347            });
348            return Err(WarmError::Embedder(EmbedderError::ModelMismatch {
349                built,
350                active,
351            }));
352        }
353        state.dim.get_or_insert(expected_dim);
354        state.built_fingerprint.get_or_insert(runtime_identity);
355        let reused_ids: Vec<String> = reused.iter().map(|(id, _)| id.clone()).collect();
356        state.vectors.extend(reused);
357        Ok(WarmOutcome {
358            reused: reused_ids,
359            missing,
360        })
361    }
362
363    /// Serialize corpus embeddings into a build-time artifact. Does not take
364    /// [`Self::operation_lock`] (does not touch cache state). An empty corpus
365    /// returns a valid zero-entry artifact without resolving the embedder.
366    pub(crate) fn build_artifact<'a, T: Embeddable + 'a>(
367        &self,
368        kind: ArtifactEntryKind,
369        items: impl IntoIterator<Item = &'a T>,
370        sink: &dyn TraceSink,
371    ) -> Result<Vec<u8>, ArtifactError> {
372        let corpus: Vec<&T> = items.into_iter().collect();
373        if corpus.is_empty() {
374            return build_empty_artifact();
375        }
376        let embedder = self.resolve_embedder(sink)?;
377        crate::embedding_artifact::build_artifact(kind, corpus, embedder.as_ref())
378    }
379
380    /// Embed any item whose id is not yet cached and insert it by id — the
381    /// incremental core of the cache. Skips ids already present, so an
382    /// already-embedded item is never recomputed (O(k) for k missing ids: newly
383    /// registered *or* invalidated-on-replace). Idempotent: a no-op once every id
384    /// is cached.
385    pub(crate) fn extend<'a, T: Embeddable + 'a>(
386        &self,
387        items: impl IntoIterator<Item = &'a T>,
388        sink: &dyn TraceSink,
389    ) -> Result<(), EmbedderError> {
390        self.with_operation_write(|cache| cache.extend_locked(items, sink))
391    }
392
393    /// Like [`Self::extend`], but assumes the caller already holds
394    /// [`Self::operation_lock`] for write.
395    pub(crate) fn extend_locked<'a, T: Embeddable + 'a>(
396        &self,
397        items: impl IntoIterator<Item = &'a T>,
398        sink: &dyn TraceSink,
399    ) -> Result<(), EmbedderError> {
400        // Gather the not-yet-cached ids so a fully-cached corpus never loads the
401        // model (empty batch → early return).
402        let missing: Vec<(String, String)> = {
403            let state = self.state.lock().expect("dense cache mutex poisoned");
404            items
405                .into_iter()
406                .filter(|item| !state.vectors.contains_key(item.embed_id()))
407                .map(|item| (item.embed_id().to_string(), item.embed_text()))
408                .collect()
409        };
410        if missing.is_empty() {
411            return Ok(());
412        }
413        let embedder = self.resolve_embedder(sink)?;
414        // One batch call: cheap for an in-process model and essential for an
415        // endpoint when an explicit build embeds the missing corpus.
416        let texts: Vec<String> = missing.iter().map(|(_, text)| text.clone()).collect();
417        let embedded = embedder.embed_batch_with_identity(&texts)?;
418        let vectors = embedded.value;
419
420        // Validate the entire batch against one staged dimension before mutating
421        // the live cache. A failure leaves all missing ids missing.
422        let existing_dim = self.state.lock().expect("dense cache mutex poisoned").dim;
423        let expected_dim =
424            validate_batch(&vectors, missing.len(), existing_dim, &embedded.fingerprint)?;
425
426        // Commit vectors and their metadata as one state transition.
427        let mut state = self.state.lock().expect("dense cache mutex poisoned");
428        if let Some(built) = &state.built_fingerprint
429            && built != &embedded.fingerprint
430        {
431            let built = built.clone();
432            let active = embedded.fingerprint;
433            sink.record(TraceEvent::EmbedderModelMismatch {
434                built: built.clone(),
435                active: active.clone(),
436            });
437            return Err(EmbedderError::ModelMismatch { built, active });
438        }
439        state.dim.get_or_insert(expected_dim);
440        state.built_fingerprint.get_or_insert(embedded.fingerprint);
441        state
442            .vectors
443            .extend(missing.into_iter().map(|(id, _)| id).zip(vectors));
444        Ok(())
445    }
446
447    /// Recompute the complete corpus, then replace vectors and metadata in one
448    /// commit. Any load, inference, identity, or dimension failure leaves the
449    /// previously searchable cache untouched.
450    pub(crate) fn rebuild<'a, T: Embeddable + 'a>(
451        &self,
452        items: impl IntoIterator<Item = &'a T>,
453        sink: &dyn TraceSink,
454    ) -> Result<(), EmbedderError> {
455        let _build = self
456            .operation_lock
457            .write()
458            .expect("dense operation lock poisoned");
459        let corpus: Vec<(String, String)> = items
460            .into_iter()
461            .map(|item| (item.embed_id().to_string(), item.embed_text()))
462            .collect();
463        if corpus.is_empty() {
464            *self.state.lock().expect("dense cache mutex poisoned") = DenseCacheState::default();
465            return Ok(());
466        }
467
468        let embedder = self.resolve_embedder(sink)?;
469        let texts: Vec<String> = corpus.iter().map(|(_, text)| text.clone()).collect();
470        let embedded = embedder.embed_batch_with_identity(&texts)?;
471        let dim = validate_batch(&embedded.value, corpus.len(), None, &embedded.fingerprint)?;
472        let vectors = corpus
473            .into_iter()
474            .map(|(id, _)| id)
475            .zip(embedded.value)
476            .collect();
477        let replacement = DenseCacheState {
478            vectors,
479            built_fingerprint: Some(embedded.fingerprint),
480            dim: Some(dim),
481        };
482        *self.state.lock().expect("dense cache mutex poisoned") = replacement;
483        Ok(())
484    }
485
486    /// Drop a cached embedding for an id whose vector must not survive.
487    ///
488    /// Two callers, for two different reasons: a registry's `register` (and the
489    /// update half of `replace_all`) drops a superseded vector so the next
490    /// [`Self::extend`] re-embeds that id; the removal half of `replace_all`
491    /// drops the vector of an id leaving the corpus entirely, with nothing to
492    /// re-embed. The second case is load-bearing rather than tidiness —
493    /// [`Self::require_built`] compares *counts*, so a vector left behind for a
494    /// departed id could offset a new, unembedded one and let the guard pass.
495    pub(crate) fn invalidate(&self, id: &str) {
496        self.state
497            .lock()
498            .expect("dense cache mutex poisoned")
499            .vectors
500            .remove(id);
501    }
502
503    /// Validate, embed, and rank one query against one immutable cache version,
504    /// returning the ranking **and** the embedded query. A concurrent
505    /// build/rebuild cannot replace the vector space between the query identity
506    /// check and cosine ranking.
507    ///
508    /// The query vector escapes because the usage-ranking arm (ADR-0014) matches
509    /// it against intent centroids: reusing it here is what makes adaptive
510    /// ranking free on the semantic and hybrid paths instead of costing a second
511    /// inference. The match happens after this returns but within the same
512    /// caller, so it observes the vector this ranking used.
513    pub(crate) fn search_returning_query_vec<'a, T: Embeddable + 'a>(
514        &self,
515        items: impl IntoIterator<Item = &'a T>,
516        query: &str,
517        depth: usize,
518        sink: &dyn TraceSink,
519    ) -> Result<RankedWithQuery, EmbedderError> {
520        let _search = self
521            .operation_lock
522            .read()
523            .expect("dense operation lock poisoned");
524        let items: Vec<&T> = items.into_iter().collect();
525        self.require_built(items.len())?;
526        let query_vec = self.embed_query(query, sink)?;
527        let ranked = self.ranked(items, &query_vec, depth);
528        Ok((ranked, query_vec))
529    }
530
531    /// Embed a query for cosine ranking (uses the same embedder as
532    /// [`Self::extend`], so the one-time model load is shared).
533    ///
534    /// Two hard guards protect against silently-wrong cosine results: a model
535    /// mismatch if the active model differs from the one that built the cache,
536    /// and a dimension mismatch if the query vector's width differs from the
537    /// corpus's (cosine over mismatched dims is meaningless, not merely worse).
538    pub(crate) fn embed_query(
539        &self,
540        query: &str,
541        sink: &dyn TraceSink,
542    ) -> Result<Vec<f32>, EmbedderError> {
543        let embedder = self.resolve_embedder(sink)?;
544
545        let built = self
546            .state
547            .lock()
548            .expect("dense cache mutex poisoned")
549            .built_fingerprint
550            .clone();
551        let embedded = embedder.embed_query_with_identity(query)?;
552        if let Some((built, active)) = model_drift(built.as_deref(), &embedded.fingerprint) {
553            sink.record(TraceEvent::EmbedderModelMismatch {
554                built: built.clone(),
555                active: active.clone(),
556            });
557            return Err(EmbedderError::ModelMismatch { built, active });
558        }
559
560        let vector = embedded.value;
561        if let Some(dim) = self.state.lock().expect("dense cache mutex poisoned").dim
562            && vector.len() != dim
563        {
564            return Err(EmbedderError::DimensionMismatch {
565                expected: dim,
566                got: vector.len(),
567                model: embedded.fingerprint,
568            });
569        }
570        Ok(vector)
571    }
572
573    /// Cosine-rank `query_vec` against the cached vectors, best-first with ties
574    /// broken by id. Assumes [`Self::extend`] already ran (`require_built`
575    /// passed), so every item's id resolves to a vector; an id missing from the
576    /// cache (shouldn't happen post-`require_built`) is simply skipped. The
577    /// id-keyed corpus holds one entry per id, so there are no duplicates to
578    /// collapse.
579    pub(crate) fn ranked<'a, T: Embeddable + 'a>(
580        &self,
581        items: impl IntoIterator<Item = &'a T>,
582        query_vec: &[f32],
583        depth: usize,
584    ) -> Vec<(String, f32)> {
585        let guard = self.state.lock().expect("dense cache mutex poisoned");
586        let docs: Vec<(String, &[f32])> = items
587            .into_iter()
588            .filter_map(|item| {
589                guard
590                    .vectors
591                    .get(item.embed_id())
592                    .map(|v| (item.embed_id().to_string(), v.as_slice()))
593            })
594            .collect();
595        dense_search(docs, query_vec, depth)
596    }
597}
598
599/// Validate one complete embedder batch without mutating cache state. Returns
600/// the common vector width that may be committed.
601fn validate_batch(
602    vectors: &[Vec<f32>],
603    expected_len: usize,
604    expected_dim: Option<usize>,
605    fingerprint: &str,
606) -> Result<usize, EmbedderError> {
607    if vectors.len() != expected_len {
608        return Err(EmbedderError::Inference {
609            source: format!(
610                "embedder returned {} embeddings for {expected_len} inputs",
611                vectors.len()
612            ),
613        });
614    }
615    let first_dim = vectors
616        .first()
617        .map(Vec::len)
618        .ok_or_else(|| EmbedderError::Inference {
619            source: "embedder returned no embeddings".into(),
620        })?;
621    let dim = expected_dim.unwrap_or(first_dim);
622    for vector in vectors {
623        if vector.len() != dim {
624            return Err(EmbedderError::DimensionMismatch {
625                expected: dim,
626                got: vector.len(),
627                model: fingerprint.to_string(),
628            });
629        }
630    }
631    Ok(dim)
632}
633
634/// Detect a model-identity mismatch: the fingerprint that built the cache vs the
635/// one now in use. `None` if the cache is unbuilt (`built` is `None`) or they match.
636/// Pure, so the drift logic is unit-tested without forcing the (currently
637/// impossible in-process) state through the cache.
638fn model_drift(built: Option<&str>, active: &str) -> Option<(String, String)> {
639    match built {
640        Some(b) if b != active => Some((b.to_string(), active.to_string())),
641        _ => None,
642    }
643}
644
645#[cfg(test)]
646mod tests {
647    use std::sync::atomic::{AtomicUsize, Ordering};
648
649    use super::*;
650    use crate::embedding::Embedded;
651    use crate::test_support::{FpCountingEmbedder, PanicOnEmbedStub, build_test_artifact, unit};
652    use crate::trace::{MemorySink, NoopSink, TraceEvent};
653
654    struct Doc {
655        id: String,
656        text: String,
657    }
658    impl Embeddable for Doc {
659        fn embed_id(&self) -> &str {
660            &self.id
661        }
662        fn embed_text(&self) -> String {
663            self.text.clone()
664        }
665    }
666    fn doc(id: &str, text: &str) -> Doc {
667        Doc {
668            id: id.into(),
669            text: text.into(),
670        }
671    }
672
673    /// One-hot embedder keyed on the word "read", counting `embed_doc` calls so
674    /// the incremental contract is provable without the network.
675    struct CountingStub {
676        docs: AtomicUsize,
677    }
678    impl CountingStub {
679        fn new() -> Self {
680            Self {
681                docs: AtomicUsize::new(0),
682            }
683        }
684        fn docs(&self) -> usize {
685            self.docs.load(Ordering::SeqCst)
686        }
687    }
688    fn vec_for(text: &str) -> Vec<f32> {
689        if text.to_lowercase().contains("read") {
690            vec![1.0, 0.0]
691        } else {
692            vec![0.0, 1.0]
693        }
694    }
695    impl Embedder for CountingStub {
696        fn embed_doc(&self, text: &str) -> Result<Vec<f32>, EmbedderError> {
697            self.docs.fetch_add(1, Ordering::SeqCst);
698            Ok(vec_for(text))
699        }
700        fn embed_query(&self, text: &str) -> Result<Vec<f32>, EmbedderError> {
701            Ok(vec_for(text))
702        }
703    }
704
705    #[test]
706    fn require_built_errors_until_the_cache_covers_the_corpus() {
707        let cache = DenseCache::with_embedder(Arc::new(CountingStub::new()));
708        let items = vec![doc("a", "read"), doc("b", "write")];
709        assert!(matches!(
710            cache.require_built(items.len()),
711            Err(EmbedderError::EmbeddingsNotBuilt)
712        ));
713        cache.extend(&items, &NoopSink).unwrap();
714        assert!(cache.require_built(items.len()).is_ok());
715    }
716
717    #[test]
718    fn extend_embeds_only_the_new_tail() {
719        let stub = Arc::new(CountingStub::new());
720        let cache = DenseCache::with_embedder(stub.clone());
721        let mut items = vec![doc("a", "read"), doc("b", "write")];
722        cache.extend(&items, &NoopSink).unwrap();
723        assert_eq!(stub.docs(), 2);
724        items.push(doc("c", "read"));
725        cache.extend(&items, &NoopSink).unwrap();
726        assert_eq!(stub.docs(), 3, "only the newly-appended item is embedded");
727        // Idempotent once caught up.
728        cache.extend(&items, &NoopSink).unwrap();
729        assert_eq!(stub.docs(), 3);
730    }
731
732    #[test]
733    fn invalidate_forces_re_embed_of_an_id() {
734        // Replace-in-place path: an id embedded as "read", then invalidated and
735        // re-embedded from "write". The new vector must win — the mechanism a
736        // re-registered tool/skill relies on (RAT-378).
737        let stub = Arc::new(CountingStub::new());
738        let cache = DenseCache::with_embedder(stub.clone());
739        cache.extend([&doc("x", "read")], &NoopSink).unwrap();
740        assert_eq!(stub.docs(), 1);
741
742        cache.invalidate("x");
743        // Re-embed x from its new content; the query matches "write".
744        cache.extend([&doc("x", "write")], &NoopSink).unwrap();
745        assert_eq!(stub.docs(), 2, "invalidated id is re-embedded, once");
746
747        let item = doc("x", "write");
748        let ranked = cache.ranked([&item], &[0.0, 1.0], 10);
749        assert_eq!(ranked.len(), 1);
750        assert_eq!(ranked[0].0, "x");
751        assert!(ranked[0].1 > 0.9, "ranks with the re-embedded vector");
752    }
753
754    /// Embeds docs at one width but the query at another — forces the query-time
755    /// dimension guard.
756    struct WidthStub {
757        doc_dim: usize,
758        query_dim: usize,
759    }
760    impl Embedder for WidthStub {
761        fn embed_doc(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
762            Ok(vec![1.0; self.doc_dim])
763        }
764        fn embed_query(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
765            Ok(vec![1.0; self.query_dim])
766        }
767    }
768
769    #[test]
770    fn query_dimension_mismatch_is_a_hard_error() {
771        let cache = DenseCache::with_embedder(Arc::new(WidthStub {
772            doc_dim: 2,
773            query_dim: 3,
774        }));
775        cache.extend([&doc("a", "x")], &NoopSink).unwrap(); // stamps dim = 2
776        let err = cache.embed_query("q", &NoopSink).unwrap_err();
777        assert!(
778            matches!(
779                err,
780                EmbedderError::DimensionMismatch {
781                    expected: 2,
782                    got: 3,
783                    ..
784                }
785            ),
786            "got: {err:?}"
787        );
788    }
789
790    struct RetryAfterMixedDimensions {
791        batches: Mutex<Vec<usize>>,
792        attempts: AtomicUsize,
793    }
794
795    impl RetryAfterMixedDimensions {
796        fn new() -> Self {
797            Self {
798                batches: Mutex::new(Vec::new()),
799                attempts: AtomicUsize::new(0),
800            }
801        }
802    }
803
804    impl Embedder for RetryAfterMixedDimensions {
805        fn embed_doc(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
806            unreachable!("test exercises the batch seam")
807        }
808
809        fn embed_query(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
810            Ok(vec![1.0, 0.0])
811        }
812
813        fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbedderError> {
814            self.batches
815                .lock()
816                .expect("batches mutex poisoned")
817                .push(texts.len());
818            if self.attempts.fetch_add(1, Ordering::SeqCst) == 0 {
819                Ok(vec![vec![1.0, 0.0], vec![1.0, 0.0, 0.0]])
820            } else {
821                Ok(vec![vec![1.0, 0.0]; texts.len()])
822            }
823        }
824    }
825
826    #[test]
827    fn failed_incremental_batch_commits_nothing_and_retries_every_missing_item() {
828        let stub = Arc::new(RetryAfterMixedDimensions::new());
829        let cache = DenseCache::with_embedder(stub.clone());
830        let items = vec![doc("a", "read"), doc("b", "write")];
831
832        assert!(matches!(
833            cache.extend(&items, &NoopSink),
834            Err(EmbedderError::DimensionMismatch { .. })
835        ));
836        cache.extend(&items, &NoopSink).unwrap();
837
838        assert_eq!(
839            *stub.batches.lock().expect("batches mutex poisoned"),
840            vec![2, 2],
841            "a failed batch must leave every item missing"
842        );
843        assert!(cache.require_built(items.len()).is_ok());
844    }
845
846    struct ChangingIdentityStub {
847        batch_identities: Mutex<std::collections::VecDeque<&'static str>>,
848        query_identity: &'static str,
849    }
850
851    impl Embedder for ChangingIdentityStub {
852        fn embed_doc(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
853            Ok(vec![1.0, 0.0])
854        }
855
856        fn embed_query(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
857            Ok(vec![1.0, 0.0])
858        }
859
860        fn embed_batch_with_identity(
861            &self,
862            texts: &[String],
863        ) -> Result<Embedded<Vec<Vec<f32>>>, EmbedderError> {
864            let fingerprint = self
865                .batch_identities
866                .lock()
867                .expect("identities mutex poisoned")
868                .pop_front()
869                .expect("scripted batch identity");
870            Ok(Embedded {
871                value: vec![vec![1.0, 0.0]; texts.len()],
872                fingerprint: fingerprint.into(),
873            })
874        }
875
876        fn embed_query_with_identity(
877            &self,
878            _text: &str,
879        ) -> Result<Embedded<Vec<f32>>, EmbedderError> {
880            Ok(Embedded {
881                value: vec![1.0, 0.0],
882                fingerprint: self.query_identity.into(),
883            })
884        }
885    }
886
887    #[test]
888    fn incremental_model_mismatch_is_hard_and_commits_nothing() {
889        let stub = Arc::new(ChangingIdentityStub {
890            batch_identities: Mutex::new(std::collections::VecDeque::from(["a", "b", "a"])),
891            query_identity: "a",
892        });
893        let cache = DenseCache::with_embedder(stub);
894        let mut items = vec![doc("a", "read")];
895        cache.extend(&items, &NoopSink).unwrap();
896        items.push(doc("b", "write"));
897
898        assert!(matches!(
899            cache.extend(&items, &NoopSink),
900            Err(EmbedderError::ModelMismatch { .. })
901        ));
902        cache.extend(&items, &NoopSink).unwrap();
903        assert!(cache.require_built(items.len()).is_ok());
904    }
905
906    #[test]
907    fn query_model_mismatch_is_a_hard_error() {
908        let stub = Arc::new(ChangingIdentityStub {
909            batch_identities: Mutex::new(std::collections::VecDeque::from(["built"])),
910            query_identity: "active",
911        });
912        let cache = DenseCache::with_embedder(stub);
913        cache.extend([&doc("a", "read")], &NoopSink).unwrap();
914
915        assert!(matches!(
916            cache.embed_query("q", &NoopSink),
917            Err(EmbedderError::ModelMismatch { built, active })
918                if built == "built" && active == "active"
919        ));
920    }
921
922    #[test]
923    fn model_drift_detects_a_changed_fingerprint() {
924        assert_eq!(model_drift(None, "a"), None, "unbuilt cache never drifts");
925        assert_eq!(model_drift(Some("a"), "a"), None, "same model never drifts");
926        assert_eq!(
927            model_drift(Some("a"), "b"),
928            Some(("a".to_string(), "b".to_string())),
929            "a changed model drifts"
930        );
931    }
932
933    #[test]
934    fn require_built_fails_after_invalidate_until_rebuilt() {
935        let cache = DenseCache::with_embedder(Arc::new(CountingStub::new()));
936        let items = vec![doc("a", "read"), doc("b", "write")];
937        cache.extend(&items, &NoopSink).unwrap();
938        assert!(cache.require_built(items.len()).is_ok());
939
940        cache.invalidate("a");
941        assert!(
942            matches!(
943                cache.require_built(items.len()),
944                Err(EmbedderError::EmbeddingsNotBuilt)
945            ),
946            "an invalidated id drops the cache below the corpus until rebuilt"
947        );
948        cache.extend(&items, &NoopSink).unwrap();
949        assert!(cache.require_built(items.len()).is_ok());
950    }
951
952    /// Warm stub with distinct runtime vs artifact identities (Local AD-1 shape).
953    struct SplitIdentityStub {
954        runtime: String,
955        artifact: String,
956    }
957
958    impl Embedder for SplitIdentityStub {
959        fn embed_doc(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
960            panic!("embed_doc must not be called during pure warm")
961        }
962        fn embed_query(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
963            panic!("embed_query must not be called during pure warm")
964        }
965        fn embed_batch(&self, _texts: &[String]) -> Result<Vec<Vec<f32>>, EmbedderError> {
966            panic!("embed_batch must not be called during pure warm")
967        }
968        fn fingerprint(&self) -> String {
969            self.runtime.clone()
970        }
971        fn artifact_identity(&self) -> Result<String, EmbedderError> {
972            Ok(self.artifact.clone())
973        }
974    }
975
976    /// Endpoint probe stub: static `fingerprint()`, probe identity via batch.
977    struct EndpointProbeStub {
978        static_fingerprint: String,
979        probe_fingerprint: String,
980        probe_calls: AtomicUsize,
981        allow_probe: bool,
982    }
983
984    impl Embedder for EndpointProbeStub {
985        fn embed_doc(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
986            panic!("embed_doc must not be called")
987        }
988        fn embed_query(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
989            panic!("embed_query must not be called")
990        }
991        fn embed_batch_with_identity(
992            &self,
993            texts: &[String],
994        ) -> Result<Embedded<Vec<Vec<f32>>>, EmbedderError> {
995            assert!(
996                self.allow_probe,
997                "Endpoint probe must not run when static fingerprint already matches"
998            );
999            assert_eq!(texts.len(), 1);
1000            self.probe_calls.fetch_add(1, Ordering::SeqCst);
1001            Ok(Embedded {
1002                value: vec![unit([1.0, 0.0])],
1003                fingerprint: self.probe_fingerprint.clone(),
1004            })
1005        }
1006        fn fingerprint(&self) -> String {
1007            self.static_fingerprint.clone()
1008        }
1009    }
1010
1011    fn sample_endpoint_model() -> EmbeddingModel {
1012        EmbeddingModel::Endpoint {
1013            url: "http://example.test/v1/embeddings".into(),
1014            model: "configured-model".into(),
1015            api_key_env: None,
1016            query_prefix: None,
1017            doc_prefix: None,
1018        }
1019    }
1020
1021    #[test]
1022    fn warm_reuses_matching_id_and_hash_without_calling_embedder() {
1023        let items = [doc("a", "read"), doc("b", "write")];
1024        let bytes = build_test_artifact(
1025            ArtifactEntryKind::Tool,
1026            &items,
1027            "fp-warm",
1028            vec![unit([1.0, 0.0]), unit([0.0, 1.0])],
1029        );
1030        let cache = DenseCache::with_embedder(Arc::new(PanicOnEmbedStub::new("fp-warm")));
1031        let outcome = cache
1032            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &items, &NoopSink)
1033            .unwrap();
1034        assert_eq!(outcome.reused, vec!["a", "b"]);
1035        assert!(outcome.missing.is_empty());
1036        assert_eq!(cache.built_fingerprint().as_deref(), Some("fp-warm"));
1037        assert_eq!(cache.dim(), Some(2));
1038        assert!(cache.require_built(items.len()).is_ok());
1039    }
1040
1041    #[test]
1042    fn warm_reports_missing_when_id_absent_from_artifact() {
1043        let artifact_items = [doc("a", "read")];
1044        let bytes = build_test_artifact(
1045            ArtifactEntryKind::Tool,
1046            &artifact_items,
1047            "fp-warm",
1048            vec![unit([1.0, 0.0])],
1049        );
1050        let corpus = [doc("a", "read"), doc("b", "write")];
1051        let cache = DenseCache::with_embedder(Arc::new(PanicOnEmbedStub::new("fp-warm")));
1052        let outcome = cache
1053            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &corpus, &NoopSink)
1054            .unwrap();
1055        assert_eq!(outcome.reused, vec!["a"]);
1056        assert_eq!(outcome.missing, vec!["b"]);
1057        assert!(cache.require_built(1).is_ok());
1058        assert!(matches!(
1059            cache.require_built(2),
1060            Err(EmbedderError::EmbeddingsNotBuilt)
1061        ));
1062    }
1063
1064    #[test]
1065    fn warm_reports_missing_when_projection_hash_differs() {
1066        let bytes = build_test_artifact(
1067            ArtifactEntryKind::Tool,
1068            &[doc("a", "read")],
1069            "fp-warm",
1070            vec![unit([1.0, 0.0])],
1071        );
1072        let corpus = [doc("a", "read changed")];
1073        let cache = DenseCache::with_embedder(Arc::new(PanicOnEmbedStub::new("fp-warm")));
1074        let outcome = cache
1075            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &corpus, &NoopSink)
1076            .unwrap();
1077        assert!(outcome.reused.is_empty());
1078        assert_eq!(outcome.missing, vec!["a"]);
1079        assert!(cache.built_fingerprint().is_none());
1080        assert!(cache.dim().is_none());
1081    }
1082
1083    #[test]
1084    fn warm_model_fingerprint_mismatch_leaves_cache_untouched() {
1085        let items = [doc("a", "read")];
1086        let bytes = build_test_artifact(
1087            ArtifactEntryKind::Tool,
1088            &items,
1089            "fp-artifact",
1090            vec![unit([1.0, 0.0])],
1091        );
1092        let cache = DenseCache::with_embedder(Arc::new(PanicOnEmbedStub::new("fp-active")));
1093        assert!(matches!(
1094            cache.warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &items, &NoopSink),
1095            Err(WarmError::ArtifactModelMismatch {
1096                artifact,
1097                active
1098            }) if artifact == "fp-artifact" && active == "fp-active"
1099        ));
1100        assert!(cache.built_fingerprint().is_none());
1101        assert!(cache.dim().is_none());
1102        assert!(matches!(
1103            cache.require_built(1),
1104            Err(EmbedderError::EmbeddingsNotBuilt)
1105        ));
1106    }
1107
1108    #[test]
1109    fn warm_ignores_other_known_entry_kind() {
1110        let items = [doc("a", "read")];
1111        let bytes = build_test_artifact(
1112            ArtifactEntryKind::Skill,
1113            &items,
1114            "fp-warm",
1115            vec![unit([1.0, 0.0])],
1116        );
1117        let cache = DenseCache::with_embedder(Arc::new(PanicOnEmbedStub::new("fp-warm")));
1118        let outcome = cache
1119            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &items, &NoopSink)
1120            .unwrap();
1121        assert!(outcome.reused.is_empty());
1122        assert_eq!(outcome.missing, vec!["a"]);
1123        assert!(cache.built_fingerprint().is_none());
1124        assert!(cache.dim().is_none());
1125    }
1126
1127    #[test]
1128    fn warm_mixed_artifact_reuses_matching_kind_only() {
1129        let tool = doc("search", "tool search text");
1130        let skill = doc("search", "skill search text");
1131        let tool_bytes = build_test_artifact(
1132            ArtifactEntryKind::Tool,
1133            std::slice::from_ref(&tool),
1134            "fp-warm",
1135            vec![unit([1.0, 0.0])],
1136        );
1137        let skill_bytes = build_test_artifact(
1138            ArtifactEntryKind::Skill,
1139            std::slice::from_ref(&skill),
1140            "fp-warm",
1141            vec![unit([0.0, 1.0])],
1142        );
1143        let bytes =
1144            crate::embedding_artifact::merge_embedding_artifacts(&[&tool_bytes, &skill_bytes])
1145                .unwrap();
1146
1147        let tool_cache = DenseCache::with_embedder(Arc::new(PanicOnEmbedStub::new("fp-warm")));
1148        let tool_outcome = tool_cache
1149            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &[tool], &NoopSink)
1150            .unwrap();
1151        assert_eq!(tool_outcome.reused, vec!["search"]);
1152        assert!(tool_outcome.missing.is_empty());
1153        assert_eq!(
1154            tool_cache.state.lock().unwrap().vectors.get("search"),
1155            Some(&unit([1.0, 0.0]))
1156        );
1157
1158        let skill_cache = DenseCache::with_embedder(Arc::new(PanicOnEmbedStub::new("fp-warm")));
1159        let skill_outcome = skill_cache
1160            .warm_from_artifact(&bytes, ArtifactEntryKind::Skill, &[skill], &NoopSink)
1161            .unwrap();
1162        assert_eq!(skill_outcome.reused, vec!["search"]);
1163        assert!(skill_outcome.missing.is_empty());
1164        assert_eq!(
1165            skill_cache.state.lock().unwrap().vectors.get("search"),
1166            Some(&unit([0.0, 1.0]))
1167        );
1168    }
1169
1170    #[test]
1171    fn warm_subset_corpus_from_superset_artifact() {
1172        let artifact_items = [doc("a", "alpha"), doc("b", "bravo"), doc("c", "charlie")];
1173        let bytes = build_test_artifact(
1174            ArtifactEntryKind::Tool,
1175            &artifact_items,
1176            "fp-warm",
1177            vec![unit([1.0, 0.0]), unit([0.0, 1.0]), unit([0.6, 0.8])],
1178        );
1179        let corpus = [doc("a", "alpha"), doc("c", "charlie")];
1180        let cache = DenseCache::with_embedder(Arc::new(PanicOnEmbedStub::new("fp-warm")));
1181        let outcome = cache
1182            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &corpus, &NoopSink)
1183            .unwrap();
1184        assert_eq!(outcome.reused, vec!["a", "c"]);
1185        assert!(outcome.missing.is_empty());
1186        assert!(cache.require_built(2).is_ok());
1187    }
1188
1189    #[test]
1190    fn warm_then_extend_embeds_only_missing_ids() {
1191        let artifact_items = [doc("a", "read file")];
1192        let bytes = build_test_artifact(
1193            ArtifactEntryKind::Tool,
1194            &artifact_items,
1195            "fp-warm",
1196            vec![unit([1.0, 0.0])],
1197        );
1198        let corpus = [doc("a", "read file"), doc("b", "write file")];
1199
1200        let count_stub = Arc::new(FpCountingEmbedder::new("fp-warm", vec_for));
1201        let cache = DenseCache::with_embedder(count_stub.clone());
1202        let outcome = cache
1203            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &corpus, &NoopSink)
1204            .unwrap();
1205        assert_eq!(outcome.reused, vec!["a"]);
1206        assert_eq!(outcome.missing, vec!["b"]);
1207        assert_eq!(count_stub.docs(), 0, "warm must not embed reused ids");
1208
1209        cache.extend(&corpus, &NoopSink).unwrap();
1210        assert_eq!(count_stub.docs(), 1, "only the missing id is embedded");
1211        assert!(cache.require_built(corpus.len()).is_ok());
1212    }
1213
1214    #[test]
1215    fn warm_on_endpoint_skips_probe_when_static_fingerprint_already_matches() {
1216        let items = [doc("a", "read")];
1217        let bytes = build_test_artifact(
1218            ArtifactEntryKind::Tool,
1219            &items,
1220            "fp-static",
1221            vec![unit([1.0, 0.0])],
1222        );
1223        let stub = Arc::new(EndpointProbeStub {
1224            static_fingerprint: "fp-static".into(),
1225            probe_fingerprint: "fp-should-not-matter".into(),
1226            probe_calls: AtomicUsize::new(0),
1227            allow_probe: false,
1228        });
1229        let cache = DenseCache::with_embedder_and_model(stub.clone(), sample_endpoint_model());
1230        let outcome = cache
1231            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &items, &NoopSink)
1232            .unwrap();
1233        assert_eq!(outcome.reused, vec!["a"]);
1234        assert_eq!(stub.probe_calls.load(Ordering::SeqCst), 0);
1235    }
1236
1237    #[test]
1238    fn warm_on_endpoint_probes_and_accepts_on_resolved_match() {
1239        let items = [doc("a", "read")];
1240        let bytes = build_test_artifact(
1241            ArtifactEntryKind::Tool,
1242            &items,
1243            "fp-resolved",
1244            vec![unit([1.0, 0.0])],
1245        );
1246        let stub = Arc::new(EndpointProbeStub {
1247            static_fingerprint: "fp-configured".into(),
1248            probe_fingerprint: "fp-resolved".into(),
1249            probe_calls: AtomicUsize::new(0),
1250            allow_probe: true,
1251        });
1252        let cache = DenseCache::with_embedder_and_model(stub.clone(), sample_endpoint_model());
1253        let outcome = cache
1254            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &items, &NoopSink)
1255            .unwrap();
1256        assert_eq!(outcome.reused, vec!["a"]);
1257        assert_eq!(stub.probe_calls.load(Ordering::SeqCst), 1);
1258        assert_eq!(cache.built_fingerprint().as_deref(), Some("fp-resolved"));
1259    }
1260
1261    #[test]
1262    fn warm_on_endpoint_probes_and_rejects_on_genuine_mismatch() {
1263        let items = [doc("a", "read")];
1264        let bytes = build_test_artifact(
1265            ArtifactEntryKind::Tool,
1266            &items,
1267            "fp-artifact",
1268            vec![unit([1.0, 0.0])],
1269        );
1270        let stub = Arc::new(EndpointProbeStub {
1271            static_fingerprint: "fp-configured".into(),
1272            probe_fingerprint: "fp-other".into(),
1273            probe_calls: AtomicUsize::new(0),
1274            allow_probe: true,
1275        });
1276        let cache = DenseCache::with_embedder_and_model(stub.clone(), sample_endpoint_model());
1277        assert!(matches!(
1278            cache.warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &items, &NoopSink),
1279            Err(WarmError::ArtifactModelMismatch {
1280                artifact,
1281                active
1282            }) if artifact == "fp-artifact" && active == "fp-other"
1283        ));
1284        assert_eq!(stub.probe_calls.load(Ordering::SeqCst), 1);
1285        assert!(cache.built_fingerprint().is_none());
1286        assert!(matches!(
1287            cache.require_built(1),
1288            Err(EmbedderError::EmbeddingsNotBuilt)
1289        ));
1290    }
1291
1292    #[test]
1293    fn warm_compares_artifact_identity_stamps_runtime() {
1294        let items = [doc("a", "read"), doc("b", "write")];
1295        let bytes = build_test_artifact(
1296            ArtifactEntryKind::Tool,
1297            &items,
1298            "content-x",
1299            vec![unit([1.0, 0.0]), unit([0.0, 1.0])],
1300        );
1301        let cache = DenseCache::with_embedder(Arc::new(SplitIdentityStub {
1302            runtime: "runtime-path-b".into(),
1303            artifact: "content-x".into(),
1304        }));
1305        let outcome = cache
1306            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &items, &NoopSink)
1307            .unwrap();
1308        assert_eq!(outcome.reused, vec!["a", "b"]);
1309        assert_eq!(
1310            cache.built_fingerprint().as_deref(),
1311            Some("runtime-path-b"),
1312            "warm must stamp runtime identity, not the RAT1 artifact identity"
1313        );
1314    }
1315
1316    #[test]
1317    fn warm_rejects_artifact_identity_mismatch() {
1318        let items = [doc("a", "read")];
1319        let bytes = build_test_artifact(
1320            ArtifactEntryKind::Tool,
1321            &items,
1322            "content-a",
1323            vec![unit([1.0, 0.0])],
1324        );
1325        let cache = DenseCache::with_embedder(Arc::new(SplitIdentityStub {
1326            runtime: "runtime-path".into(),
1327            artifact: "content-b".into(),
1328        }));
1329        assert!(matches!(
1330            cache.warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &items, &NoopSink),
1331            Err(WarmError::ArtifactModelMismatch {
1332                artifact,
1333                active
1334            }) if artifact == "content-a" && active == "content-b"
1335        ));
1336        assert!(cache.built_fingerprint().is_none());
1337        assert!(cache.dim().is_none());
1338    }
1339
1340    #[test]
1341    fn artifact_model_mismatch_message_names_the_artifact() {
1342        let items = [doc("a", "read")];
1343        let bytes = build_test_artifact(
1344            ArtifactEntryKind::Tool,
1345            &items,
1346            "content-a",
1347            vec![unit([1.0, 0.0])],
1348        );
1349        let cache = DenseCache::with_embedder(Arc::new(SplitIdentityStub {
1350            runtime: "runtime-path".into(),
1351            artifact: "content-b".into(),
1352        }));
1353        let err = cache
1354            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &items, &NoopSink)
1355            .expect_err("header mismatch");
1356        let message = err.to_string();
1357        assert!(
1358            message.contains("embedding artifact was built with"),
1359            "{message}"
1360        );
1361        assert!(message.contains("rebuild the artifact"), "{message}");
1362        assert!(!message.contains("cache was built with"), "{message}");
1363        assert!(!message.contains("re-embed the corpus"), "{message}");
1364    }
1365
1366    #[test]
1367    fn warm_rejects_nonempty_zero_dim_before_cache_mutation() {
1368        let items = [doc("a", "read")];
1369        let bytes = crate::embedding_artifact::test_hand_artifact(
1370            crate::embedding_artifact::projection_version(),
1371            0,
1372            "fp-zero-dim",
1373            &[ArtifactEntry {
1374                kind: ArtifactEntryKind::Tool,
1375                id: "a".into(),
1376                projection_hash: hash_projection_text("read"),
1377                vector: vec![],
1378            }],
1379        );
1380        let cache = DenseCache::with_embedder(Arc::new(PanicOnEmbedStub::new("unused")));
1381        assert!(matches!(
1382            cache.warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &items, &NoopSink),
1383            Err(WarmError::Artifact(ArtifactError::NonEmptyZeroDim))
1384        ));
1385        assert!(cache.built_fingerprint().is_none());
1386        assert!(cache.dim().is_none());
1387    }
1388
1389    #[test]
1390    fn warm_into_prebuilt_rejects_dimension_mismatch() {
1391        let cache = DenseCache::with_embedder(Arc::new(WidthStub {
1392            doc_dim: 3,
1393            query_dim: 3,
1394        }));
1395        let prebuilt = doc("a", "x");
1396        cache.extend([&prebuilt], &NoopSink).unwrap();
1397        assert_eq!(cache.dim(), Some(3));
1398        assert_eq!(cache.built_fingerprint().as_deref(), Some("unknown"));
1399        let a_before = cache
1400            .state
1401            .lock()
1402            .unwrap()
1403            .vectors
1404            .get("a")
1405            .cloned()
1406            .expect("prebuilt id a");
1407
1408        let warm_items = [doc("b", "y")];
1409        let bytes = build_test_artifact(
1410            ArtifactEntryKind::Tool,
1411            &warm_items,
1412            "unknown",
1413            vec![unit([1.0, 0.0])],
1414        );
1415        assert!(matches!(
1416            cache.warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &warm_items, &NoopSink),
1417            Err(WarmError::Embedder(EmbedderError::DimensionMismatch {
1418                expected: 3,
1419                got: 2,
1420                ..
1421            }))
1422        ));
1423        assert_eq!(cache.dim(), Some(3));
1424        assert_eq!(cache.built_fingerprint().as_deref(), Some("unknown"));
1425        let state = cache.state.lock().unwrap();
1426        assert_eq!(state.vectors.get("a"), Some(&a_before));
1427        assert!(!state.vectors.contains_key("b"));
1428        assert_eq!(state.vectors.len(), 1);
1429    }
1430
1431    #[test]
1432    fn warm_into_prebuilt_rejects_runtime_fingerprint_mismatch() {
1433        let stub = Arc::new(EndpointProbeStub {
1434            static_fingerprint: "fp-static-Y".into(),
1435            probe_fingerprint: "fp-probed-X".into(),
1436            probe_calls: AtomicUsize::new(0),
1437            allow_probe: true,
1438        });
1439        let cache = DenseCache::with_embedder_and_model(stub.clone(), sample_endpoint_model());
1440        let prebuilt = doc("a", "read");
1441        cache.extend([&prebuilt], &NoopSink).unwrap();
1442        assert_eq!(cache.built_fingerprint().as_deref(), Some("fp-probed-X"));
1443        assert_eq!(cache.dim(), Some(2));
1444        let a_before = cache
1445            .state
1446            .lock()
1447            .unwrap()
1448            .vectors
1449            .get("a")
1450            .cloned()
1451            .expect("prebuilt id a");
1452        let probes_before = stub.probe_calls.load(Ordering::SeqCst);
1453        assert_eq!(probes_before, 1);
1454
1455        let warm_items = [doc("b", "write")];
1456        let bytes = build_test_artifact(
1457            ArtifactEntryKind::Tool,
1458            &warm_items,
1459            "fp-static-Y",
1460            vec![unit([0.0, 1.0])],
1461        );
1462        let sink = MemorySink::new("warm-second-guard");
1463        let err = cache
1464            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &warm_items, &sink)
1465            .expect_err("runtime fingerprint mismatch");
1466        assert!(matches!(
1467            &err,
1468            WarmError::Embedder(EmbedderError::ModelMismatch {
1469                built,
1470                active
1471            }) if built == "fp-probed-X" && active == "fp-static-Y"
1472        ));
1473        assert!(err.to_string().contains("cache was built with"), "{err}");
1474        assert_eq!(stub.probe_calls.load(Ordering::SeqCst), probes_before);
1475        let mismatch: Vec<_> = sink
1476            .drain()
1477            .into_iter()
1478            .filter_map(|e| match e.event {
1479                TraceEvent::EmbedderModelMismatch { built, active } => Some((built, active)),
1480                _ => None,
1481            })
1482            .collect();
1483        assert_eq!(
1484            mismatch,
1485            vec![("fp-probed-X".into(), "fp-static-Y".into())],
1486            "second guard must emit exactly one EmbedderModelMismatch"
1487        );
1488        assert_eq!(cache.dim(), Some(2));
1489        assert_eq!(cache.built_fingerprint().as_deref(), Some("fp-probed-X"));
1490        let state = cache.state.lock().unwrap();
1491        assert_eq!(state.vectors.get("a"), Some(&a_before));
1492        assert!(!state.vectors.contains_key("b"));
1493        assert_eq!(state.vectors.len(), 1);
1494    }
1495
1496    #[test]
1497    fn warm_into_prebuilt_adds_compatible_id() {
1498        let stub = Arc::new(FpCountingEmbedder::new("fp-add", vec_for));
1499        let cache = DenseCache::with_embedder(stub.clone());
1500        let prebuilt = doc("a", "read");
1501        cache.extend([&prebuilt], &NoopSink).unwrap();
1502        assert_eq!(stub.docs(), 1);
1503        assert_eq!(cache.dim(), Some(2));
1504        assert_eq!(cache.built_fingerprint().as_deref(), Some("fp-add"));
1505
1506        let warm_items = [doc("b", "write")];
1507        let bytes = build_test_artifact(
1508            ArtifactEntryKind::Tool,
1509            &warm_items,
1510            "fp-add",
1511            vec![unit([0.0, 1.0])],
1512        );
1513        let docs_before_warm = stub.docs();
1514        let outcome = cache
1515            .warm_from_artifact(&bytes, ArtifactEntryKind::Tool, &warm_items, &NoopSink)
1516            .unwrap();
1517        assert_eq!(outcome.reused, vec!["b"]);
1518        assert!(outcome.missing.is_empty());
1519        assert_eq!(stub.docs(), docs_before_warm);
1520        assert_eq!(cache.dim(), Some(2));
1521        assert_eq!(cache.built_fingerprint().as_deref(), Some("fp-add"));
1522        let state = cache.state.lock().unwrap();
1523        assert!(state.vectors.contains_key("a"));
1524        assert!(state.vectors.contains_key("b"));
1525        assert_eq!(state.vectors.len(), 2);
1526    }
1527}