//! Tests for the étage-2 vector recall (LOT C, §14.5.3): //! - [`should_use_vector`] — the pure routing decision (full truth table); //! - [`AdaptiveMemoryRecall`] — the étage-1/étage-2 switch, with naïve equivalence //! on strategy `none`, routing by size/budget, and the never-fail-hard fallback; //! - [`VectorMemoryRecall`] — semantic ranking with the deterministic //! [`HashEmbedder`], budget truncation, the derived `vectors.json` cache and its //! embedder-id/dimension invalidation, and `StubEmbedder` best-effort behaviour; //! - [`HashEmbedder`] / [`StubEmbedder`] embedder contracts; //! - [`embedder_from_profile`] strategy → embedder mapping. //! //! Everything is in-memory and deterministic (`MemFs` + `HashEmbedder`); no real //! I/O, no network. use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use async_trait::async_trait; use domain::markdown::MarkdownDoc; use domain::memory::{Memory, MemoryFrontmatter, MemoryIndexEntry, MemorySlug, MemoryType}; use domain::ports::{ DirEntry, Embedder, EmbedderError, FileSystem, FsError, MemoryQuery, MemoryRecall, MemoryStore, RemotePath, }; use domain::profile::{EmbedderProfile, EmbedderStrategy}; use domain::project::ProjectPath; use infrastructure::{ embedder_from_profile, should_use_vector, AdaptiveMemoryRecall, FsMemoryStore, HashEmbedder, NaiveMemoryRecall, StubEmbedder, VectorMemoryRecall, }; // --------------------------------------------------------------------------- // In-memory FileSystem mock (same shape as `memory_recall.rs` / `memory_store.rs`), // with raw read access to assert on the derived `vectors.json`. // --------------------------------------------------------------------------- #[derive(Default)] struct MemFs { files: Mutex>>, } impl MemFs { fn arc() -> Arc { Arc::new(Self::default()) } fn raw(&self, path: &str) -> Option { self.files .lock() .unwrap() .get(path) .map(|b| String::from_utf8(b.clone()).unwrap()) } } #[async_trait] impl FileSystem for MemFs { async fn read(&self, path: &RemotePath) -> Result, FsError> { self.files .lock() .unwrap() .get(path.as_str()) .cloned() .ok_or_else(|| FsError::NotFound(path.as_str().to_string())) } async fn write(&self, path: &RemotePath, data: &[u8]) -> Result<(), FsError> { self.files .lock() .unwrap() .insert(path.as_str().to_string(), data.to_vec()); Ok(()) } async fn exists(&self, path: &RemotePath) -> Result { Ok(self.files.lock().unwrap().contains_key(path.as_str())) } async fn create_dir_all(&self, _path: &RemotePath) -> Result<(), FsError> { Ok(()) } async fn list(&self, _path: &RemotePath) -> Result, FsError> { Ok(Vec::new()) } async fn symlink(&self, _src: &RemotePath, _dst: &RemotePath) -> Result<(), FsError> { Ok(()) } } const VECTORS_PATH: &str = "/proj/.ideai/memory/.index/vectors.json"; // --------------------------------------------------------------------------- // A counting wrapper embedder: delegates to an inner embedder and counts the // number of *texts* it was asked to embed — to assert the derived cache avoids // recomputing note vectors on a second recall. // --------------------------------------------------------------------------- struct SpyEmbedder { inner: HashEmbedder, embedded_texts: AtomicUsize, } impl SpyEmbedder { fn new(id: &str, dim: usize) -> Self { Self { inner: HashEmbedder::new(id, dim), embedded_texts: AtomicUsize::new(0), } } } #[async_trait] impl Embedder for SpyEmbedder { fn id(&self) -> &str { self.inner.id() } async fn embed(&self, texts: &[String]) -> Result>, EmbedderError> { self.embedded_texts.fetch_add(texts.len(), Ordering::SeqCst); self.inner.embed(texts).await } fn dimension(&self) -> usize { self.inner.dimension() } } // --------------------------------------------------------------------------- // Builders. // --------------------------------------------------------------------------- fn root() -> ProjectPath { ProjectPath::new("/proj").unwrap() } /// A note whose index line is ` ` (title == slug). The vector recall /// embeds exactly `"{slug} {hook}"`, so picking distinct vocabulary words in /// `slug`/`hook` makes cosine similarity to a chosen query deterministic. fn note(slug: &str, hook: &str) -> Memory { Memory::new( MemoryFrontmatter { name: MemorySlug::new(slug).unwrap(), description: hook.to_string(), r#type: MemoryType::Project, }, MarkdownDoc::new("# body"), ) .unwrap() } fn query(text: &str, budget: usize) -> MemoryQuery { MemoryQuery { text: text.to_string(), token_budget: budget, } } fn slugs(entries: &[MemoryIndexEntry]) -> Vec { entries.iter().map(|e| e.slug.as_str().to_string()).collect() } async fn store_with(notes: &[Memory], fs: Arc) -> Arc { let store = Arc::new(FsMemoryStore::new(fs)); for n in notes { store.save(&root(), n).await.unwrap(); } store } // =========================================================================== // 1. should_use_vector — pure truth table. // =========================================================================== #[test] fn should_use_vector_none_strategy_is_always_false() { // Even when size > budget, a `none` strategy never routes to the vector path. assert!(!should_use_vector(0, 0, EmbedderStrategy::None)); assert!(!should_use_vector(100, 10, EmbedderStrategy::None)); assert!(!should_use_vector(10, 100, EmbedderStrategy::None)); assert!(!should_use_vector(usize::MAX, 0, EmbedderStrategy::None)); } #[test] fn should_use_vector_non_none_iff_size_exceeds_budget() { for strategy in [ EmbedderStrategy::LocalOnnx, EmbedderStrategy::LocalServer, EmbedderStrategy::Api, ] { // size > budget ⇒ true. assert!(should_use_vector(11, 10, strategy), "{strategy:?}: 11>10"); assert!(should_use_vector(1, 0, strategy), "{strategy:?}: 1>0"); // size == budget ⇒ false (exact boundary). assert!(!should_use_vector(10, 10, strategy), "{strategy:?}: 10==10"); assert!(!should_use_vector(0, 0, strategy), "{strategy:?}: 0==0"); // size < budget ⇒ false. assert!(!should_use_vector(9, 10, strategy), "{strategy:?}: 9<10"); } } // =========================================================================== // 2. HashEmbedder — determinism, dimension, L2 normalisation. // =========================================================================== #[tokio::test] async fn hash_embedder_is_deterministic() { let e = HashEmbedder::new("h", 64); let a = e.embed(&["hello world".to_string()]).await.unwrap(); let b = e.embed(&["hello world".to_string()]).await.unwrap(); assert_eq!(a, b, "same text must give the same vector"); } #[tokio::test] async fn hash_embedder_dimension_is_consistent() { let e = HashEmbedder::new("h", 48); assert_eq!(e.dimension(), 48); let out = e .embed(&["alpha beta".to_string(), "gamma".to_string()]) .await .unwrap(); assert_eq!(out.len(), 2, "one vector per input, order preserved"); assert!(out.iter().all(|v| v.len() == 48), "every vector is dimension-long"); } #[tokio::test] async fn hash_embedder_vectors_are_l2_normalised() { let e = HashEmbedder::new("h", 64); let out = e .embed(&["one two three four".to_string()]) .await .unwrap(); let norm = out[0].iter().map(|x| x * x).sum::().sqrt(); assert!((norm - 1.0).abs() < 1e-5, "norm ≈ 1, got {norm}"); } #[tokio::test] async fn hash_embedder_empty_text_yields_zero_vector_no_panic() { // No tokens ⇒ zero vector (norm 0, left unnormalised). Must not panic / NaN. let e = HashEmbedder::new("h", 16); let out = e.embed(&["".to_string()]).await.unwrap(); assert_eq!(out[0].len(), 16); assert!(out[0].iter().all(|x| *x == 0.0)); } // =========================================================================== // 3. StubEmbedder — always Unsupported, never panics. // =========================================================================== #[tokio::test] async fn stub_embedder_returns_unsupported() { let e = StubEmbedder::new("stub", 32, "localOnnx"); assert_eq!(e.dimension(), 32); let err = e.embed(&["x".to_string()]).await.unwrap_err(); assert!(matches!(err, EmbedderError::Unsupported(_)), "got {err:?}"); } // =========================================================================== // 4. embedder_from_profile — strategy → embedder mapping. // =========================================================================== #[tokio::test] async fn embedder_from_profile_maps_each_strategy() { // none ⇒ no embedder (recall stays naïve). assert!(embedder_from_profile(&EmbedderProfile::none()).is_none()); // localOnnx / localServer / api ⇒ a StubEmbedder (Unsupported), dimension kept. for strategy in [ EmbedderStrategy::LocalOnnx, EmbedderStrategy::LocalServer, EmbedderStrategy::Api, ] { let profile = EmbedderProfile::new("e", "E", strategy, None, None, None, 24).unwrap(); let embedder = embedder_from_profile(&profile) .unwrap_or_else(|| panic!("{strategy:?} must yield an embedder")); assert_eq!(embedder.dimension(), 24); let err = embedder.embed(&["t".to_string()]).await.unwrap_err(); assert!( matches!(err, EmbedderError::Unsupported(_)), "{strategy:?} stub must be Unsupported, got {err:?}" ); } } // =========================================================================== // 5. VectorMemoryRecall — semantic ranking, budget, cache. // =========================================================================== /// Builds a vector recall over `notes` with the given embedder, sharing one fs. fn vector_recall( embedder: Arc, store: Arc, fs: Arc, ) -> VectorMemoryRecall { VectorMemoryRecall::new(embedder, store, fs) } #[tokio::test] async fn vector_recall_ranks_closest_note_first() { // Notes have disjoint vocabularies; the query shares all tokens with `beta`, // so `beta` must rank first (cosine == 1 for it, ~0 for the others). let fs = MemFs::arc(); let notes = [ note("alpha", "apple orange grape"), note("beta", "kiwi mango papaya"), note("gamma", "carrot potato onion"), ]; let store = store_with(¬es, fs.clone()).await; let embedder = Arc::new(HashEmbedder::new("hash", 256)); let recall = vector_recall(embedder, store, fs); // Query identical to beta's index line ("beta kiwi mango papaya"). let out = recall .recall(&root(), &query("beta kiwi mango papaya", 100_000)) .await .unwrap(); assert_eq!(out.first().map(|e| e.slug.as_str()), Some("beta"), "ranked: {:?}", slugs(&out)); // All three notes fit the ample budget. assert_eq!(out.len(), 3); } #[tokio::test] async fn vector_recall_truncates_to_budget() { // The closest note ("beta apricot apricot") must survive a tight budget while // the rest are pruned. Cost of an entry = ceil((title+hook chars)/4). let fs = MemFs::arc(); let notes = [ note("beta", "kiwi"), // "beta kiwi" → 8 chars → 2 tokens note("alpha", "apple orange grape"), // larger note("gamma", "carrot potato onion"), // larger ]; let store = store_with(¬es, fs.clone()).await; let embedder = Arc::new(HashEmbedder::new("hash", 256)); let recall = vector_recall(embedder, store, fs); // Budget 2 fits only the top-ranked `beta` (cost 2); the next would exceed. let out = recall.recall(&root(), &query("beta kiwi", 2)).await.unwrap(); assert_eq!(slugs(&out), vec!["beta"], "budget must keep only the closest"); } #[tokio::test] async fn vector_recall_zero_budget_and_empty_memory_are_empty() { let fs = MemFs::arc(); let embedder = Arc::new(HashEmbedder::new("hash", 64)); // Empty memory ⇒ empty, no error. let empty_store = Arc::new(FsMemoryStore::new(fs.clone())); let recall = vector_recall(embedder.clone(), empty_store, fs.clone()); assert!(recall .recall(&root(), &query("anything", 1000)) .await .unwrap() .is_empty()); // Zero budget ⇒ empty, even with notes present. let store = store_with(&[note("a", "x"), note("b", "y")], fs.clone()).await; let recall = vector_recall(embedder, store, fs); assert!(recall .recall(&root(), &query("a x", 0)) .await .unwrap() .is_empty()); } #[tokio::test] async fn vector_recall_writes_and_reuses_the_derived_cache() { let fs = MemFs::arc(); let notes = [note("alpha", "apple"), note("beta", "banana")]; let store = store_with(¬es, fs.clone()).await; let spy = Arc::new(SpyEmbedder::new("spy", 128)); let recall = VectorMemoryRecall::new(spy.clone(), store, fs.clone()); // First recall: embeds 2 note texts + 1 query text = 3. let _ = recall.recall(&root(), &query("alpha apple", 100_000)).await.unwrap(); let after_first = spy.embedded_texts.load(Ordering::SeqCst); assert_eq!(after_first, 3, "2 notes + 1 query embedded on a cold cache"); // The derived cache file now exists with both slugs and the embedder id. let doc = fs.raw(VECTORS_PATH).expect("vectors.json written"); assert!(doc.contains("\"embedderId\": \"spy\""), "tagged with embedder id: {doc}"); assert!(doc.contains("\"alpha\""), "alpha vector cached: {doc}"); assert!(doc.contains("\"beta\""), "beta vector cached: {doc}"); // Second recall: note vectors come from the cache; only the query is embedded. let _ = recall.recall(&root(), &query("beta banana", 100_000)).await.unwrap(); let delta = spy.embedded_texts.load(Ordering::SeqCst) - after_first; assert_eq!(delta, 1, "cache reuse ⇒ only the query is re-embedded, got {delta}"); } #[tokio::test] async fn vector_recall_invalidates_cache_on_embedder_id_change() { let fs = MemFs::arc(); let notes = [note("alpha", "apple"), note("beta", "banana")]; let store = store_with(¬es, fs.clone()).await; // Warm the cache with embedder "first". let first = Arc::new(SpyEmbedder::new("first", 128)); VectorMemoryRecall::new(first.clone(), store.clone(), fs.clone()) .recall(&root(), &query("alpha apple", 100_000)) .await .unwrap(); assert!(fs.raw(VECTORS_PATH).unwrap().contains("\"embedderId\": \"first\"")); // A different embedder id must invalidate ⇒ recompute all note vectors. let second = Arc::new(SpyEmbedder::new("second", 128)); VectorMemoryRecall::new(second.clone(), store, fs.clone()) .recall(&root(), &query("beta banana", 100_000)) .await .unwrap(); // 2 notes recomputed + 1 query embedded against the new embedder. assert_eq!( second.embedded_texts.load(Ordering::SeqCst), 3, "id change must invalidate and recompute note vectors" ); // The file is now tagged with the new embedder id. assert!(fs.raw(VECTORS_PATH).unwrap().contains("\"embedderId\": \"second\"")); } #[tokio::test] async fn vector_recall_invalidates_cache_on_dimension_change() { let fs = MemFs::arc(); let notes = [note("alpha", "apple"), note("beta", "banana")]; let store = store_with(¬es, fs.clone()).await; // Warm the cache at dimension 64. let dim64 = Arc::new(SpyEmbedder::new("same-id", 64)); VectorMemoryRecall::new(dim64, store.clone(), fs.clone()) .recall(&root(), &query("alpha apple", 100_000)) .await .unwrap(); // Same id, different dimension ⇒ invalidation, all note vectors recomputed. let dim128 = Arc::new(SpyEmbedder::new("same-id", 128)); VectorMemoryRecall::new(dim128.clone(), store, fs.clone()) .recall(&root(), &query("beta banana", 100_000)) .await .unwrap(); assert_eq!( dim128.embedded_texts.load(Ordering::SeqCst), 3, "dimension change must invalidate and recompute" ); assert!(fs.raw(VECTORS_PATH).unwrap().contains("\"dimension\": 128")); } #[tokio::test] async fn vector_recall_with_stub_embedder_errors_without_panic() { // The vector recall surfaces the embedder error as a MemoryError (the Adaptive // switch turns that into a naïve fallback — covered below). Here we just pin // that StubEmbedder is best-effort: an Err, never a panic. let fs = MemFs::arc(); let store = store_with(&[note("alpha", "apple")], fs.clone()).await; let stub = Arc::new(StubEmbedder::new("stub", 64, "localOnnx")); let recall = VectorMemoryRecall::new(stub, store, fs); let res = recall.recall(&root(), &query("alpha apple", 1000)).await; assert!(res.is_err(), "stub embedder ⇒ Err on the vector path, got {res:?}"); } // =========================================================================== // 6. AdaptiveMemoryRecall — the étage-1/étage-2 switch. // =========================================================================== /// Assembles an [`AdaptiveMemoryRecall`] over `notes` for the given strategy and /// embedder, sharing one in-memory fs / store. async fn adaptive_over( notes: &[Memory], strategy: EmbedderStrategy, embedder: Arc, ) -> (Arc, AdaptiveMemoryRecall) { let fs = MemFs::arc(); let store = store_with(notes, fs.clone()).await; let naive: Arc = Arc::new(NaiveMemoryRecall::new(store.clone())); let vector: Arc = Arc::new(VectorMemoryRecall::new(embedder, store.clone(), fs.clone())); let adaptive = AdaptiveMemoryRecall::new(naive, vector, store, strategy); (fs, adaptive) } #[tokio::test] async fn adaptive_none_strategy_matches_naive_exactly() { let notes = [ note("alpha", "apple orange grape"), note("beta", "kiwi mango papaya"), note("gamma", "carrot potato onion"), ]; let fs = MemFs::arc(); let store = store_with(¬es, fs.clone()).await; let naive = NaiveMemoryRecall::new(store.clone()); // Adaptive with strategy `none` (embedder present but never used). let naive_dyn: Arc = Arc::new(NaiveMemoryRecall::new(store.clone())); let vector_dyn: Arc = Arc::new(VectorMemoryRecall::new( Arc::new(HashEmbedder::new("hash", 64)), store.clone(), fs.clone(), )); let adaptive = AdaptiveMemoryRecall::new(naive_dyn, vector_dyn, store, EmbedderStrategy::None); // For several budgets, adaptive(none) == naive, byte-for-byte. for budget in [0usize, 3, 6, 100, 100_000] { let q = query("kiwi mango papaya", budget); let expected = naive.recall(&root(), &q).await.unwrap(); let got = adaptive.recall(&root(), &q).await.unwrap(); assert_eq!(got, expected, "strategy none must equal naive at budget {budget}"); } } #[tokio::test] async fn adaptive_small_memory_routes_to_naive_order() { // Memory smaller than the budget ⇒ naïve path ⇒ index order, even with a // non-none strategy. The query would, on the vector path, reorder by // similarity; index order proves naïve was used. let notes = [ note("alpha", "apple"), note("beta", "banana"), note("gamma", "grape"), ]; let (_fs, adaptive) = adaptive_over( ¬es, EmbedderStrategy::LocalServer, Arc::new(HashEmbedder::new("hash", 256)), ) .await; // Huge budget ⇒ memory_size <= budget ⇒ naïve ⇒ index order alpha,beta,gamma. let out = adaptive .recall(&root(), &query("gamma grape", 100_000)) .await .unwrap(); assert_eq!(slugs(&out), vec!["alpha", "beta", "gamma"], "naïve keeps index order"); } #[tokio::test] async fn adaptive_large_memory_routes_to_vector_ranking() { // Memory larger than the budget + non-none strategy ⇒ vector path ⇒ the note // closest to the query ranks first (not index order). let notes = [ note("alpha", "apple orange grape"), note("beta", "kiwi mango papaya"), note("gamma", "carrot potato onion"), ]; let (_fs, adaptive) = adaptive_over( ¬es, EmbedderStrategy::LocalOnnx, Arc::new(HashEmbedder::new("hash", 256)), ) .await; // Budget smaller than the total memory size forces the vector route; the query // matches `gamma`, so `gamma` must be first (≠ index order, where it is last). let total = { let store = FsMemoryStore::new(MemFs::arc()); for n in ¬es { store.save(&root(), n).await.unwrap(); } infrastructure::index_token_size(&store.read_index(&root()).await.unwrap()) }; assert!(total > 4, "sanity: memory must exceed the chosen budget"); let out = adaptive .recall(&root(), &query("carrot potato onion", total - 1)) .await .unwrap(); assert_eq!( out.first().map(|e| e.slug.as_str()), Some("gamma"), "vector path ranks the closest note first: {:?}", slugs(&out) ); } #[tokio::test] async fn adaptive_falls_back_to_naive_when_embedder_fails() { // Large memory + non-none strategy routes to the vector path; with a // StubEmbedder that path errors, and Adaptive must degrade to naïve — never a // hard error. Result therefore equals the naïve recall (index order, truncated). let notes = [ note("alpha", "apple orange grape"), note("beta", "kiwi mango papaya"), note("gamma", "carrot potato onion"), ]; let fs = MemFs::arc(); let store = store_with(¬es, fs.clone()).await; let naive_ref = NaiveMemoryRecall::new(store.clone()); let naive_dyn: Arc = Arc::new(NaiveMemoryRecall::new(store.clone())); let vector_dyn: Arc = Arc::new(VectorMemoryRecall::new( Arc::new(StubEmbedder::new("stub", 64, "localOnnx")), store.clone(), fs.clone(), )); let adaptive = AdaptiveMemoryRecall::new(naive_dyn, vector_dyn, store.clone(), EmbedderStrategy::Api); let total = infrastructure::index_token_size(&store.read_index(&root()).await.unwrap()); let q = query("carrot potato onion", total - 1); let got = adaptive.recall(&root(), &q).await.expect("fallback must not error"); let expected = naive_ref.recall(&root(), &q).await.unwrap(); assert_eq!(got, expected, "embedder failure must degrade to the naïve result"); assert!(!got.is_empty(), "the degraded result still returns entries"); } #[tokio::test] async fn adaptive_empty_memory_and_zero_budget_are_empty() { // Liskov common contract: empty list, never an error, both paths. let fs = MemFs::arc(); let empty_store = Arc::new(FsMemoryStore::new(fs.clone())); let naive: Arc = Arc::new(NaiveMemoryRecall::new(empty_store.clone())); let vector: Arc = Arc::new(VectorMemoryRecall::new( Arc::new(HashEmbedder::new("hash", 64)), empty_store.clone(), fs.clone(), )); let adaptive = AdaptiveMemoryRecall::new(naive, vector, empty_store, EmbedderStrategy::LocalOnnx); // Empty memory. assert!(adaptive .recall(&root(), &query("anything", 1000)) .await .unwrap() .is_empty()); // Zero budget. assert!(adaptive .recall(&root(), &query("anything", 0)) .await .unwrap() .is_empty()); }