apend

Author: zTgx

src/client/builder.rs

@@ -507,6 +507,17 @@ impl EngineBuilder {
                 retriever.with_content_config(retrieval_config.content.to_aggregator_config());
         }
 
+        // Add memo store if provided or create default
+        if let Some(memo_store) = self.memo_store {
+            retriever = retriever.with_memo_store(memo_store);
+        } else {
+            // Create default memo store with model from config
+            let memo_store = MemoStore::new()
+                .with_model(&retrieval_config.model)
+                .with_version(1);
+            retriever = retriever.with_memo_store(memo_store);
+        }
+
         // Build engine
         Engine::with_components(config, workspace, retriever, executor)
             .await

src/memo/store.rs

@@ -42,6 +42,44 @@ struct MemoStoreData {
     stats: MemoStats,
 }
 
+/// Atomic statistics for lock-free access.
+#[derive(Debug, Default)]
+struct AtomicStats {
+    hits: AtomicU64,
+    misses: AtomicU64,
+    tokens_saved: AtomicU64,
+}
+
+impl AtomicStats {
+    fn new() -> Self {
+        Self {
+            hits: AtomicU64::new(0),
+            misses: AtomicU64::new(0),
+            tokens_saved: AtomicU64::new(0),
+        }
+    }
+
+    fn record_hit(&self) {
+        self.hits.fetch_add(1, Ordering::Relaxed);
+    }
+
+    fn record_miss(&self) {
+        self.misses.fetch_add(1, Ordering::Relaxed);
+    }
+
+    fn add_tokens_saved(&self, tokens: u64) {
+        self.tokens_saved.fetch_add(tokens, Ordering::Relaxed);
+    }
+
+    fn snapshot(&self) -> (u64, u64, u64) {
+        (
+            self.hits.load(Ordering::Relaxed),
+            self.misses.load(Ordering::Relaxed),
+            self.tokens_saved.load(Ordering::Relaxed),
+        )
+    }
+}
+
 /// LLM Memoization store.
 ///
 /// Provides caching for expensive LLM operations with:
@@ -79,6 +117,17 @@ pub struct MemoStore {
     version: u32,
 }
 
+impl std::fmt::Debug for MemoStore {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("MemoStore")
+            .field("ttl", &self.ttl)
+            .field("model_id", &self.model_id)
+            .field("version", &self.version)
+            .field("cache_len", &self.cache.read().len())
+            .finish()
+    }
+}
+
 impl Clone for MemoStore {
     fn clone(&self) -> Self {
         Self {
@@ -639,8 +688,13 @@ mod tests {
             .await
             .unwrap();
 
-        // Hit
-        store.get(&key);
+        // Hit via get_or_compute (this updates global stats)
+        store
+            .get_or_compute(key.clone(), || async {
+                Ok((MemoValue::Summary("Should not be called".to_string()), 0))
+            })
+            .await
+            .unwrap();
 
         let stats = store.stats().await;
         assert_eq!(stats.misses, 1);

src/retrieval/pipeline_retriever.rs

@@ -200,6 +200,7 @@ impl Clone for PipelineRetriever {
             max_backtracks: self.max_backtracks,
             max_iterations: self.max_iterations,
             content_config: self.content_config.clone(),
+            memo_store: self.memo_store.clone(),
         }
     }
 }