Skip to main content

context_forge/
engine.rs

1//! Core business logic: assembly, scoring, and snapshot management.
2
3use std::sync::Arc;
4use std::time::{SystemTime, UNIX_EPOCH};
5use uuid::Uuid;
6
7use crate::config::{Config, DEFAULT_RECENCY_HALF_LIFE_SECS};
8use crate::entry::ContextEntry;
9use crate::error::Error;
10use crate::lexicon::LexiconScorer;
11use crate::traits::{ContextStorage, Result, Searcher};
12// Bring the Embedder trait into scope for method dispatch on Arc<FasEmbedder>.
13#[cfg(feature = "semantic")]
14use crate::semantic::Embedder as _;
15
16/// Default candidate limit when fetching search results for assembly.
17const DEFAULT_SEARCH_LIMIT: usize = 50;
18
19/// Well-defined query constant that means "match all entries" in the Searcher trait.
20///
21/// FTS5 interprets `*` as a prefix wildcard that matches every token.
22/// Searcher implementations MUST treat this value as a match-all query.
23pub const MATCH_ALL_QUERY: &str = "*";
24
25/// Options for saving a snapshot entry.
26#[derive(Debug, Default, Clone)]
27pub struct SaveOptions {
28    /// Optional caller session identifier.
29    pub session_id: Option<String>,
30    /// Namespace partition for the new entry. `None` = global scope.
31    pub scope: Option<String>,
32    /// Arbitrary caller metadata, stored as JSON.
33    ///
34    /// Unlike `content`, metadata is persisted **verbatim** by
35    /// [`crate::ContextForge::save`] — it is not passed through
36    /// [`crate::scrub_secrets`]. Callers MUST NOT place untrusted or
37    /// secret-bearing text in this field.
38    pub metadata: Option<serde_json::Value>,
39}
40
41/// Estimate token count using whitespace heuristic (1 token ≈ 4 chars).
42#[must_use]
43pub fn estimate_tokens(text: &str) -> usize {
44    text.len().div_ceil(4)
45}
46
47/// Exponential recency decay.
48///
49/// Returns a multiplier in `(0.0, 1.0]` where 1.0 means "just created" and
50/// the value halves every `half_life` seconds.
51fn recency_decay(age_seconds: f64, half_life: f64) -> f64 {
52    0.5_f64.powf(age_seconds / half_life)
53}
54
55/// The core context engine that coordinates storage, search, and assembly.
56pub struct ContextEngine {
57    storage: Box<dyn ContextStorage>,
58    searcher: Box<dyn Searcher>,
59    config: Config,
60    scorer: Option<Arc<dyn LexiconScorer>>,
61    #[cfg(feature = "semantic")]
62    embedder: Option<Arc<crate::semantic::FasEmbedder>>,
63}
64
65impl ContextEngine {
66    /// Create a new engine with the given storage backend, searcher, and config.
67    ///
68    /// If `config.recency_half_life_secs` is not positive and finite, it is
69    /// clamped to [`DEFAULT_RECENCY_HALF_LIFE_SECS`] to prevent NaN/inf in
70    /// recency decay scoring.
71    #[must_use]
72    pub fn new(
73        storage: Box<dyn ContextStorage>,
74        searcher: Box<dyn Searcher>,
75        mut config: Config,
76    ) -> Self {
77        if !config.recency_half_life_secs.is_finite() || config.recency_half_life_secs <= 0.0 {
78            config.recency_half_life_secs = DEFAULT_RECENCY_HALF_LIFE_SECS;
79        }
80
81        Self {
82            storage,
83            searcher,
84            config,
85            scorer: None,
86            #[cfg(feature = "semantic")]
87            embedder: None,
88        }
89    }
90
91    /// Attach a [`LexiconScorer`] to this engine.
92    ///
93    /// The scorer runs inside [`Self::assemble`] after BM25 + recency decay
94    /// and before the token-budget cut. Calling this again replaces the previous
95    /// scorer. Prefer constructing via [`crate::ContextForgeBuilder`], which
96    /// automatically composes [`crate::DefaultEnglishScorer`] with any
97    /// caller-provided persona scorer.
98    #[must_use]
99    pub fn with_scorer(mut self, scorer: Arc<dyn LexiconScorer>) -> Self {
100        self.scorer = Some(scorer);
101        self
102    }
103
104    /// Attach an [`crate::semantic::Embedder`] for semantic search.
105    ///
106    /// When set, [`Self::save_snapshot`] generates and stores an embedding for
107    /// each new entry, and [`Self::assemble`] blends BM25 and semantic
108    /// candidates via Reciprocal Rank Fusion (RRF, k=60).
109    #[cfg(feature = "semantic")]
110    #[must_use]
111    pub(crate) fn with_embedder(mut self, embedder: Arc<crate::semantic::FasEmbedder>) -> Self {
112        self.embedder = Some(embedder);
113        self
114    }
115
116    /// Assemble context entries that fit within `token_budget`.
117    ///
118    /// 1. Searches for candidates matching `query`, restricted to `scope` if given.
119    /// 2. Applies recency weighting to each candidate's score.
120    /// 3. Sorts by weighted score descending.
121    /// 4. Packs entries greedily until the budget is exhausted.
122    ///
123    /// `scope = None` searches every entry regardless of scope (global recall).
124    /// `scope = Some(s)` restricts the search to entries whose `scope` equals `s`.
125    ///
126    /// # Errors
127    ///
128    /// Returns an error if the underlying search fails.
129    pub async fn assemble(
130        &self,
131        query: &str,
132        scope: Option<&str>,
133        token_budget: usize,
134    ) -> Result<Vec<ContextEntry>> {
135        let limit = DEFAULT_SEARCH_LIMIT;
136
137        tracing::debug!(query = %query, ?scope, token_budget, "assemble: bm25 search");
138        let bm25_candidates = self.searcher.search(query, scope, limit).await?;
139        tracing::debug!(count = %bm25_candidates.len(), "assemble: bm25 complete");
140
141        // Semantic search runs when the feature is enabled and an embedder is
142        // configured; otherwise returns empty (no-op default on the Searcher trait).
143        let semantic_candidates = self.run_semantic_search(query, scope, limit).await;
144        tracing::debug!(count = %semantic_candidates.len(), "assemble: semantic complete");
145
146        if bm25_candidates.is_empty() && semantic_candidates.is_empty() {
147            tracing::debug!("assemble: both searches empty");
148            return Ok(Vec::new());
149        }
150
151        // Reciprocal Rank Fusion (RRF, k=60) over the full union of both
152        // candidate sets. Entries absent from one list get worst_rank = limit+1.
153        //
154        // RRF score = 1/(k + bm25_rank) + 1/(k + semantic_rank)
155        // Then multiply by recency decay and (1 + lexicon_boost).
156        let worst_rank = limit + 1;
157        let k = 60.0_f64;
158        let now = current_timestamp();
159        let half_life = self.config.recency_half_life_secs;
160
161        // Build entry map: semantic first so BM25 entries win on overlap.
162        let mut entry_map: std::collections::HashMap<String, ContextEntry> =
163            std::collections::HashMap::new();
164        for se in &semantic_candidates {
165            entry_map
166                .entry(se.entry.id.clone())
167                .or_insert_with(|| se.entry.clone());
168        }
169        for se in &bm25_candidates {
170            entry_map.insert(se.entry.id.clone(), se.entry.clone());
171        }
172
173        let bm25_rank: std::collections::HashMap<&str, usize> = bm25_candidates
174            .iter()
175            .enumerate()
176            .map(|(i, se)| (se.entry.id.as_str(), i + 1))
177            .collect();
178        let semantic_rank: std::collections::HashMap<&str, usize> = semantic_candidates
179            .iter()
180            .enumerate()
181            .map(|(i, se)| (se.entry.id.as_str(), i + 1))
182            .collect();
183
184        let mut weighted: Vec<(f64, ContextEntry)> = entry_map
185            .into_iter()
186            .map(|(id, entry)| {
187                let br = *bm25_rank.get(id.as_str()).unwrap_or(&worst_rank);
188                let sr = *semantic_rank.get(id.as_str()).unwrap_or(&worst_rank);
189                #[allow(
190                    clippy::cast_precision_loss,
191                    reason = "ranks are small integers (max ~51); lossless"
192                )]
193                let rrf = 1.0 / (k + br as f64) + 1.0 / (k + sr as f64);
194
195                #[allow(
196                    clippy::cast_precision_loss,
197                    reason = "Unix timestamps fit losslessly in f64 for millions of years"
198                )]
199                let age = (now - entry.timestamp).max(0) as f64;
200                let decay = recency_decay(age, half_life);
201                let boost = self
202                    .scorer
203                    .as_ref()
204                    .map_or(0.0_f32, |s| s.score(&entry, query));
205                let final_score = rrf * decay * (1.0 + f64::from(boost).clamp(-1.0, 2.0));
206                (final_score, entry)
207            })
208            .collect();
209
210        weighted.sort_by(|a, b| b.0.total_cmp(&a.0));
211        tracing::debug!(fused = %weighted.len(), "assemble: rrf fusion complete");
212
213        // Greedy bin-packing by token budget — skip oversized entries, don't abort.
214        let mut result = Vec::new();
215        let mut tokens_used: usize = 0;
216        for (_score, entry) in weighted {
217            let entry_tokens = entry
218                .token_count
219                .unwrap_or_else(|| estimate_tokens(&entry.content));
220            if tokens_used.saturating_add(entry_tokens) > token_budget {
221                continue;
222            }
223            tokens_used += entry_tokens;
224            result.push(entry);
225        }
226
227        tracing::debug!(entries = %result.len(), tokens_used, "assemble: complete");
228        Ok(result)
229    }
230
231    /// Run semantic search if an embedder is available; otherwise return empty.
232    #[allow(
233        clippy::unused_async,
234        reason = "async only active under semantic feature"
235    )]
236    async fn run_semantic_search(
237        &self,
238        query: &str,
239        scope: Option<&str>,
240        limit: usize,
241    ) -> Vec<crate::entry::ScoredEntry> {
242        #[cfg(feature = "semantic")]
243        if let Some(ref embedder) = self.embedder {
244            let emb = embedder.clone();
245            let query_owned = query.to_owned();
246            match tokio::task::spawn_blocking(move || emb.embed(&query_owned)).await {
247                Ok(Ok(embedding)) => {
248                    tracing::debug!(dims = %embedding.len(), "query embedding ready");
249                    match self
250                        .searcher
251                        .search_semantic(&embedding, scope, limit)
252                        .await
253                    {
254                        Ok(results) => return results,
255                        Err(e) => tracing::warn!(error = %e, "semantic search failed"),
256                    }
257                }
258                Ok(Err(e)) => tracing::warn!(error = %e, "query embedding failed"),
259                Err(e) => tracing::warn!(error = %e, "embed task panicked"),
260            }
261        }
262        // Suppress unused-parameter lint when semantic feature is disabled.
263        #[cfg(not(feature = "semantic"))]
264        let _ = (query, scope, limit);
265        Vec::new()
266    }
267
268    /// Save a new snapshot entry. Capacity enforcement (LRU eviction) is
269    /// handled atomically by the storage layer inside an immediate
270    /// transaction, so no in-process locking is required here.
271    ///
272    /// Returns the generated entry ID.
273    ///
274    /// # Security
275    ///
276    /// This is a low-level write path that does **not** scrub secrets from
277    /// `content`. [`ContextForge::save`](crate::ContextForge::save) is the
278    /// only entry point that applies [`scrub_secrets`](crate::scrub_secrets)
279    /// before persistence; callers writing through the engine directly are
280    /// responsible for scrubbing first.
281    ///
282    /// # Errors
283    ///
284    /// Returns [`Error::InvalidEntry`] if `content` is empty, or propagates
285    /// any error from the underlying storage write.
286    pub async fn save_snapshot(
287        &self,
288        content: &str,
289        kind: &str,
290        options: &SaveOptions,
291    ) -> Result<String> {
292        if content.is_empty() {
293            return Err(Error::InvalidEntry("content must not be empty".into()));
294        }
295
296        let timestamp = current_timestamp();
297        let id = Uuid::now_v7().to_string();
298        let token_count = estimate_tokens(content);
299
300        let entry = ContextEntry {
301            id: id.clone(),
302            content: content.to_owned(),
303            timestamp,
304            kind: kind.to_owned(),
305            scope: options.scope.clone(),
306            session_id: options.session_id.clone(),
307            token_count: Some(token_count),
308            metadata: options.metadata.clone(),
309        };
310
311        self.storage.save(&entry).await?;
312        tracing::trace!(id = %id, kind = %kind, tokens = %token_count, "entry saved");
313
314        // Non-fatal: if embedding fails, the entry is stored and BM25-searchable.
315        self.embed_and_store(&id, content).await;
316
317        Ok(id)
318    }
319
320    /// Generate and persist an embedding for an already-saved entry.
321    ///
322    /// Errors are logged and swallowed — the entry is always stored even when
323    /// embedding fails.
324    #[allow(
325        clippy::unused_async,
326        reason = "async only active under semantic feature"
327    )]
328    async fn embed_and_store(&self, id: &str, content: &str) {
329        #[cfg(feature = "semantic")]
330        if let Some(ref embedder) = self.embedder {
331            tracing::debug!(id = %id, "generating embedding");
332            let emb = embedder.clone();
333            let text = content.to_owned();
334            let id_owned = id.to_owned();
335            match tokio::task::spawn_blocking(move || emb.embed(&text)).await {
336                Ok(Ok(embedding)) => {
337                    if let Err(e) = self.storage.save_embedding(&id_owned, &embedding).await {
338                        tracing::warn!(
339                            id = %id_owned,
340                            error = %e,
341                            "save_embedding failed; entry stored without semantic index"
342                        );
343                    } else {
344                        tracing::trace!(id = %id_owned, dims = %embedding.len(), "embedding stored");
345                    }
346                }
347                Ok(Err(e)) => tracing::warn!(
348                    id = %id,
349                    error = %e,
350                    "embed failed; entry stored without semantic index"
351                ),
352                Err(e) => tracing::warn!(
353                    id = %id,
354                    error = %e,
355                    "embed task panicked; entry stored without semantic index"
356                ),
357            }
358        }
359        #[cfg(not(feature = "semantic"))]
360        let _ = (id, content);
361    }
362
363    /// Return a reference to the underlying storage backend.
364    ///
365    /// This is exposed for callers that need direct access to storage
366    /// operations (delete, clear, count) not covered by [`Self::assemble`]
367    /// or [`Self::save_snapshot`].
368    #[must_use]
369    pub fn storage(&self) -> &dyn ContextStorage {
370        self.storage.as_ref()
371    }
372
373    /// Backfill embeddings for all entries that do not yet have one.
374    ///
375    /// Entries are fetched and embedded in batches of `batch_size` (capped to
376    /// at least 1). The `progress` callback is called after each batch with
377    /// `(done_so_far, total)`.
378    ///
379    /// Returns the number of entries that were successfully embedded.
380    ///
381    /// # Errors
382    ///
383    /// Returns an error if fetching unembedded entries fails or if the
384    /// embedding task panics.
385    #[cfg(feature = "semantic")]
386    pub async fn backfill_embeddings(
387        &self,
388        batch_size: usize,
389        progress: impl Fn(usize, usize),
390    ) -> Result<usize> {
391        let embedder = if let Some(e) = &self.embedder {
392            e.clone()
393        } else {
394            tracing::debug!("backfill_embeddings: no embedder configured");
395            return Ok(0);
396        };
397
398        let all = self.storage.get_unembedded(usize::MAX).await?;
399        let total = all.len();
400        if total == 0 {
401            tracing::debug!("backfill_embeddings: all entries already embedded");
402            return Ok(0);
403        }
404
405        tracing::debug!(total = %total, batch_size = %batch_size, "backfill_embeddings: starting");
406        let batch = batch_size.max(1);
407        let mut n_embedded = 0usize;
408
409        for chunk in all.chunks(batch) {
410            let texts: Vec<String> = chunk.iter().map(|e| e.content.clone()).collect();
411            let emb = embedder.clone();
412            let embeddings = tokio::task::spawn_blocking(move || {
413                let refs: Vec<&str> = texts.iter().map(String::as_str).collect();
414                emb.embed_batch(&refs)
415            })
416            .await
417            .map_err(|e| Error::Migration(format!("backfill embed task panicked: {e}")))?
418            .map_err(|e| Error::Migration(format!("backfill embed failed: {e}")))?;
419
420            for (entry, embedding) in chunk.iter().zip(embeddings.iter()) {
421                if let Err(e) = self.storage.save_embedding(&entry.id, embedding).await {
422                    tracing::warn!(id = %entry.id, error = %e, "backfill: save_embedding failed");
423                } else {
424                    n_embedded += 1;
425                }
426            }
427
428            progress(n_embedded, total);
429            tracing::debug!(done = %n_embedded, total = %total, "backfill_embeddings: batch done");
430        }
431
432        Ok(n_embedded)
433    }
434}
435
436/// Current Unix timestamp in seconds.
437///
438/// Returns `0` if the system clock reports a time before the Unix epoch,
439/// which should not happen on any supported platform.
440fn current_timestamp() -> i64 {
441    let secs = SystemTime::now()
442        .duration_since(UNIX_EPOCH)
443        .map_or(0, |d| d.as_secs());
444    i64::try_from(secs).unwrap_or(i64::MAX)
445}
446
447#[cfg(test)]
448mod tests {
449    use super::*;
450    use crate::config::{EvictionPolicy, DEFAULT_RECENCY_HALF_LIFE_SECS};
451    use crate::entry::ScoredEntry;
452    use async_trait::async_trait;
453    use std::path::PathBuf;
454    use std::sync::Mutex;
455
456    // -----------------------------------------------------------------------
457    // Mock implementations
458    // -----------------------------------------------------------------------
459
460    struct MockStorage {
461        entries: Mutex<Vec<ContextEntry>>,
462    }
463
464    impl MockStorage {
465        fn new() -> Self {
466            Self {
467                entries: Mutex::new(Vec::new()),
468            }
469        }
470    }
471
472    #[async_trait]
473    impl ContextStorage for MockStorage {
474        async fn save(&self, entry: &ContextEntry) -> Result<()> {
475            self.entries.lock().unwrap().push(entry.clone());
476            Ok(())
477        }
478
479        async fn get_top_k(&self, k: usize) -> Result<Vec<ContextEntry>> {
480            let guard = self.entries.lock().unwrap();
481            let mut sorted = guard.clone();
482            sorted.sort_by_key(|e| std::cmp::Reverse(e.timestamp));
483            sorted.truncate(k);
484            Ok(sorted)
485        }
486
487        async fn get_all(&self) -> Result<Vec<ContextEntry>> {
488            Ok(self.entries.lock().unwrap().clone())
489        }
490
491        async fn delete(&self, id: &str) -> Result<bool> {
492            let mut guard = self.entries.lock().unwrap();
493            let before = guard.len();
494            guard.retain(|e| e.id != id);
495            Ok(guard.len() < before)
496        }
497
498        async fn clear(&self) -> Result<usize> {
499            let mut guard = self.entries.lock().unwrap();
500            let n = guard.len();
501            guard.clear();
502            Ok(n)
503        }
504
505        async fn clear_scope(&self, scope: &str) -> Result<usize> {
506            let mut guard = self.entries.lock().unwrap();
507            let before = guard.len();
508            guard.retain(|e| e.scope.as_deref() != Some(scope));
509            Ok(before - guard.len())
510        }
511
512        async fn count(&self) -> Result<usize> {
513            Ok(self.entries.lock().unwrap().len())
514        }
515    }
516
517    struct MockSearcher {
518        results: Mutex<Vec<ScoredEntry>>,
519    }
520
521    impl MockSearcher {
522        fn new(results: Vec<ScoredEntry>) -> Self {
523            Self {
524                results: Mutex::new(results),
525            }
526        }
527
528        fn empty() -> Self {
529            Self::new(Vec::new())
530        }
531    }
532
533    #[async_trait]
534    impl Searcher for MockSearcher {
535        async fn search(
536            &self,
537            _query: &str,
538            _scope: Option<&str>,
539            limit: usize,
540        ) -> Result<Vec<ScoredEntry>> {
541            let guard = self.results.lock().unwrap();
542            Ok(guard.iter().take(limit).cloned().collect())
543        }
544    }
545
546    fn default_config(max_entries: usize) -> Config {
547        Config {
548            max_entries,
549            token_budget: 8192,
550            db_path: PathBuf::from(":memory:"),
551            eviction_policy: EvictionPolicy::Lru,
552            recency_half_life_secs: DEFAULT_RECENCY_HALF_LIFE_SECS,
553            ..Config::default()
554        }
555    }
556
557    fn make_entry(id: &str, content: &str, timestamp: i64) -> ContextEntry {
558        ContextEntry {
559            id: id.into(),
560            content: content.into(),
561            timestamp,
562            kind: crate::entry::kind::MANUAL.to_owned(),
563            scope: None,
564            session_id: None,
565            token_count: Some(estimate_tokens(content)),
566            metadata: None,
567        }
568    }
569
570    fn make_scored(id: &str, content: &str, timestamp: i64, score: f64) -> ScoredEntry {
571        ScoredEntry {
572            entry: make_entry(id, content, timestamp),
573            score,
574        }
575    }
576
577    // -----------------------------------------------------------------------
578    // Tests
579    // -----------------------------------------------------------------------
580
581    #[test]
582    fn test_estimate_tokens() {
583        assert_eq!(estimate_tokens(""), 0);
584        assert_eq!(estimate_tokens("a"), 1); // (1+3)/4 = 1
585        assert_eq!(estimate_tokens("ab"), 1); // (2+3)/4 = 1
586        assert_eq!(estimate_tokens("abc"), 1); // (3+3)/4 = 1 (integer)
587        assert_eq!(estimate_tokens("abcd"), 1); // (4+3)/4 = 1
588        assert_eq!(estimate_tokens("abcde"), 2); // (5+3)/4 = 2
589        assert_eq!(estimate_tokens("abcdefgh"), 2); // (8+3)/4 = 2 (integer)
590        assert_eq!(estimate_tokens("hello world, this is a test"), 7);
591    }
592
593    #[test]
594    fn test_engine_send_sync() {
595        fn assert_send_sync<T: Send + Sync>() {}
596        assert_send_sync::<ContextEngine>();
597    }
598
599    #[tokio::test]
600    async fn test_assemble_fits_budget() {
601        let now = current_timestamp();
602        let results = vec![
603            make_scored("a", "short", now, 1.0),
604            make_scored("b", "medium length text here", now, 0.9),
605            make_scored("c", "another entry with more content inside", now, 0.8),
606        ];
607
608        let engine = ContextEngine::new(
609            Box::new(MockStorage::new()),
610            Box::new(MockSearcher::new(results)),
611            default_config(100),
612        );
613
614        // Budget of 2 tokens: "short" = 2 tokens fits exactly.
615        let assembled = engine.assemble("test", None, 2).await.unwrap();
616        assert_eq!(assembled.len(), 1);
617        assert_eq!(assembled[0].id, "a");
618
619        // Budget large enough for all.
620        let assembled = engine.assemble("test", None, 1000).await.unwrap();
621        assert_eq!(assembled.len(), 3);
622    }
623
624    #[tokio::test]
625    async fn test_assemble_skips_oversized_entries() {
626        // B1 regression: oversized top-ranked entry must not abort the loop.
627        let now = current_timestamp();
628        let big_content = "x".repeat(4000); // 1000 tokens
629        let results = vec![
630            make_scored("big", &big_content, now, 1.0),
631            make_scored("small", "fits", now, 0.9),
632        ];
633
634        let engine = ContextEngine::new(
635            Box::new(MockStorage::new()),
636            Box::new(MockSearcher::new(results)),
637            default_config(100),
638        );
639
640        // Budget = 5: "big" (1000 tokens) is skipped, "fits" (1 token) is returned.
641        let assembled = engine.assemble("test", None, 5).await.unwrap();
642        assert_eq!(assembled.len(), 1);
643        assert_eq!(assembled[0].id, "small");
644    }
645
646    #[tokio::test]
647    async fn test_assemble_empty_results() {
648        let engine = ContextEngine::new(
649            Box::new(MockStorage::new()),
650            Box::new(MockSearcher::empty()),
651            default_config(100),
652        );
653        let assembled = engine.assemble("anything", None, 1000).await.unwrap();
654        assert!(assembled.is_empty());
655    }
656
657    #[tokio::test]
658    async fn test_recency_weighting() {
659        let now = current_timestamp();
660        // Two entries with equal BM25 scores but different timestamps.
661        let results = vec![
662            make_scored("old", "old entry", now - 86_400, 1.0), // 24h old
663            make_scored("new", "new entry", now, 1.0),          // just now
664        ];
665
666        let engine = ContextEngine::new(
667            Box::new(MockStorage::new()),
668            Box::new(MockSearcher::new(results)),
669            default_config(100),
670        );
671
672        let assembled = engine.assemble("test", None, 1000).await.unwrap();
673        assert_eq!(assembled.len(), 2);
674        // Newer entry should rank first due to higher recency score.
675        assert_eq!(assembled[0].id, "new");
676        assert_eq!(assembled[1].id, "old");
677    }
678
679    #[tokio::test]
680    async fn test_save_snapshot_creates_entry() {
681        let engine = ContextEngine::new(
682            Box::new(MockStorage::new()),
683            Box::new(MockSearcher::empty()),
684            default_config(100),
685        );
686
687        let id = engine
688            .save_snapshot(
689                "hello world",
690                crate::entry::kind::MANUAL,
691                &SaveOptions::default(),
692            )
693            .await
694            .unwrap();
695        assert!(!id.is_empty());
696
697        // Verify the entry exists in storage.
698        let all = engine.storage.get_all().await.unwrap();
699        assert_eq!(all.len(), 1);
700        assert_eq!(all[0].id, id);
701        assert_eq!(all[0].content, "hello world");
702        assert_eq!(all[0].token_count, Some(estimate_tokens("hello world")));
703        assert_eq!(all[0].kind, crate::entry::kind::MANUAL);
704    }
705
706    #[tokio::test]
707    async fn test_save_snapshot_populates_scope_and_metadata() {
708        let engine = ContextEngine::new(
709            Box::new(MockStorage::new()),
710            Box::new(MockSearcher::empty()),
711            default_config(100),
712        );
713
714        let metadata = serde_json::json!({"source": "test"});
715        let options = SaveOptions {
716            session_id: Some("sess-1".to_owned()),
717            scope: Some("project:homelab-rs".to_owned()),
718            metadata: Some(metadata.clone()),
719        };
720
721        let id = engine
722            .save_snapshot("hello scoped", crate::entry::kind::SNAPSHOT, &options)
723            .await
724            .unwrap();
725
726        let all = engine.storage.get_all().await.unwrap();
727        assert_eq!(all.len(), 1);
728        assert_eq!(all[0].id, id);
729        assert_eq!(all[0].kind, crate::entry::kind::SNAPSHOT);
730        assert_eq!(all[0].scope.as_deref(), Some("project:homelab-rs"));
731        assert_eq!(all[0].session_id.as_deref(), Some("sess-1"));
732        assert_eq!(all[0].metadata, Some(metadata));
733    }
734
735    #[test]
736    fn test_recency_decay_values() {
737        let half_life = crate::config::DEFAULT_RECENCY_HALF_LIFE_HOURS * 3600.0;
738
739        // At age 0, decay should be 1.0.
740        let decay_0 = recency_decay(0.0, half_life);
741        assert!((decay_0 - 1.0).abs() < f64::EPSILON);
742
743        // At age = half_life, decay should be 0.5.
744        let decay_half = recency_decay(half_life, half_life);
745        assert!((decay_half - 0.5).abs() < 1e-10);
746
747        // At age = 2 * half_life, decay should be 0.25.
748        let decay_double = recency_decay(2.0 * half_life, half_life);
749        assert!((decay_double - 0.25).abs() < 1e-10);
750    }
751
752    #[tokio::test]
753    async fn test_save_snapshot_ids_are_unique() {
754        let engine = ContextEngine::new(
755            Box::new(MockStorage::new()),
756            Box::new(MockSearcher::empty()),
757            default_config(100),
758        );
759
760        let id1 = engine
761            .save_snapshot(
762                "identical content",
763                crate::entry::kind::SNAPSHOT,
764                &SaveOptions::default(),
765            )
766            .await
767            .unwrap();
768        let id2 = engine
769            .save_snapshot(
770                "identical content",
771                crate::entry::kind::SNAPSHOT,
772                &SaveOptions::default(),
773            )
774            .await
775            .unwrap();
776        assert_ne!(
777            id1, id2,
778            "Two saves of identical content must produce distinct IDs"
779        );
780    }
781
782    #[tokio::test]
783    async fn test_save_empty_content_rejected() {
784        let engine = ContextEngine::new(
785            Box::new(MockStorage::new()),
786            Box::new(MockSearcher::empty()),
787            default_config(100),
788        );
789
790        let result = engine
791            .save_snapshot("", crate::entry::kind::MANUAL, &SaveOptions::default())
792            .await;
793        assert!(result.is_err());
794        let err = result.unwrap_err();
795        assert!(err.to_string().contains("content must not be empty"));
796    }
797
798    #[test]
799    fn test_invalid_half_life_clamped_to_default() {
800        let invalid_values: Vec<f64> = vec![
801            0.0,
802            -1.0,
803            -100.0,
804            f64::NAN,
805            f64::INFINITY,
806            f64::NEG_INFINITY,
807        ];
808
809        for value in invalid_values {
810            let mut config = default_config(100);
811            config.recency_half_life_secs = value;
812
813            let engine = ContextEngine::new(
814                Box::new(MockStorage::new()),
815                Box::new(MockSearcher::empty()),
816                config,
817            );
818
819            assert!(
820                (engine.config.recency_half_life_secs - DEFAULT_RECENCY_HALF_LIFE_SECS).abs()
821                    < f64::EPSILON,
822                "half_life {value} should have been clamped to default",
823            );
824        }
825    }
826
827    #[tokio::test]
828    async fn lexicon_scorer_boosts_important_entry_over_neutral() {
829        use crate::lexicon::DefaultEnglishScorer;
830
831        let now = current_timestamp();
832        // Equal BM25 scores, equal timestamps — only differentiator is lexicon signal.
833        let results = vec![
834            make_scored("neutral", "something unrelated here", now, 1.0),
835            make_scored("important", "confirmed, that is correct", now, 1.0),
836        ];
837
838        let scorer: Arc<dyn LexiconScorer> = Arc::new(DefaultEnglishScorer::default());
839        let engine = ContextEngine::new(
840            Box::new(MockStorage::new()),
841            Box::new(MockSearcher::new(results)),
842            default_config(100),
843        )
844        .with_scorer(scorer);
845
846        let assembled = engine.assemble("test", None, 1000).await.unwrap();
847        assert_eq!(assembled.len(), 2);
848        assert_eq!(
849            assembled[0].id, "important",
850            "lexicon-boosted entry should rank first"
851        );
852    }
853
854    #[test]
855    fn test_valid_half_life_preserved() {
856        let mut config = default_config(100);
857        config.recency_half_life_secs = 3600.0; // 1 hour — valid
858
859        let engine = ContextEngine::new(
860            Box::new(MockStorage::new()),
861            Box::new(MockSearcher::empty()),
862            config,
863        );
864
865        assert!((engine.config.recency_half_life_secs - 3600.0).abs() < f64::EPSILON);
866    }
867}