1#![allow(missing_docs)]
4use std::collections::{HashMap, HashSet};
5use std::future::Future;
6use std::pin::Pin;
7
8use parking_lot::RwLock;
9
10use chrono::Utc;
11
12use crate::auth::TenantScope;
13use crate::error::Error;
14
15use super::bm25;
16use super::hybrid;
17use super::scoring::{STRENGTH_DECAY_RATE, ScoringWeights, composite_score, effective_strength};
18use super::{Memory, MemoryEntry, MemoryQuery};
19
20pub const IN_MEMORY_STORE_DEFAULT_CAP: usize = 100_000;
27
28#[derive(Debug, Clone, Default)]
33struct EntryTokens {
34 lower_content: String,
38 content_words: Vec<String>,
42 lower_keywords: Vec<String>,
45}
46
47fn build_entry_tokens(entry: &MemoryEntry) -> EntryTokens {
48 let lower_content = entry.content.to_lowercase();
49 let content_words: Vec<String> = lower_content.split_whitespace().map(String::from).collect();
50 let lower_keywords: Vec<String> = entry.keywords.iter().map(|k| k.to_lowercase()).collect();
51 EntryTokens {
52 lower_content,
53 content_words,
54 lower_keywords,
55 }
56}
57
58fn index_entry(
62 inverted: &mut HashMap<String, HashSet<String>>,
63 entry_id: &str,
64 tokens: &EntryTokens,
65) {
66 let mut seen: HashSet<&str> = HashSet::with_capacity(tokens.content_words.len());
67 for word in tokens
68 .content_words
69 .iter()
70 .chain(tokens.lower_keywords.iter())
71 {
72 if seen.insert(word.as_str()) {
73 inverted
74 .entry(word.clone())
75 .or_default()
76 .insert(entry_id.to_string());
77 }
78 }
79}
80
81fn deindex_entry(
86 inverted: &mut HashMap<String, HashSet<String>>,
87 entry_id: &str,
88 tokens: &EntryTokens,
89) {
90 let mut seen: HashSet<&str> = HashSet::with_capacity(tokens.content_words.len());
91 for word in tokens
92 .content_words
93 .iter()
94 .chain(tokens.lower_keywords.iter())
95 {
96 if !seen.insert(word.as_str()) {
97 continue;
98 }
99 if let Some(bucket) = inverted.get_mut(word) {
100 bucket.remove(entry_id);
101 if bucket.is_empty() {
102 inverted.remove(word);
103 }
104 }
105 }
106}
107
108pub struct InMemoryStore {
130 entries: RwLock<HashMap<String, MemoryEntry>>,
131 tokens: RwLock<HashMap<String, EntryTokens>>,
132 inverted: RwLock<HashMap<String, HashSet<String>>>,
137 scoring_weights: ScoringWeights,
138 max_entries: usize,
139}
140
141impl InMemoryStore {
142 pub fn new() -> Self {
143 Self {
144 entries: RwLock::new(HashMap::new()),
145 tokens: RwLock::new(HashMap::new()),
146 inverted: RwLock::new(HashMap::new()),
147 scoring_weights: ScoringWeights::default(),
148 max_entries: IN_MEMORY_STORE_DEFAULT_CAP,
149 }
150 }
151
152 pub fn with_scoring_weights(mut self, weights: ScoringWeights) -> Self {
153 self.scoring_weights = weights;
154 self
155 }
156
157 pub fn with_max_entries(mut self, max_entries: usize) -> Self {
159 self.max_entries = max_entries;
160 self
161 }
162}
163
164impl Default for InMemoryStore {
165 fn default() -> Self {
166 Self::new()
167 }
168}
169
170impl Memory for InMemoryStore {
171 fn store(
172 &self,
173 scope: &TenantScope,
174 mut entry: MemoryEntry,
175 ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
176 entry.author_tenant_id = Some(scope.tenant_id.clone());
178 entry.author_user_id = scope.user_id.clone();
179 Box::pin(async move {
180 let mut entries = self.entries.write();
181 let mut tokens = self.tokens.write();
182 let mut inverted = self.inverted.write();
183 if !entries.contains_key(&entry.id) && entries.len() >= self.max_entries {
190 let now = Utc::now();
191 if let Some(victim_id) = entries
192 .values()
193 .min_by(|a, b| {
194 let ea = effective_strength(
195 a.strength,
196 a.last_accessed,
197 now,
198 STRENGTH_DECAY_RATE,
199 );
200 let eb = effective_strength(
201 b.strength,
202 b.last_accessed,
203 now,
204 STRENGTH_DECAY_RATE,
205 );
206 ea.partial_cmp(&eb).unwrap_or(std::cmp::Ordering::Equal)
207 })
208 .map(|e| e.id.clone())
209 {
210 entries.remove(&victim_id);
211 if let Some(victim_tokens) = tokens.remove(&victim_id) {
212 deindex_entry(&mut inverted, &victim_id, &victim_tokens);
213 }
214 tracing::warn!(
215 evicted = %victim_id,
216 cap = self.max_entries,
217 "InMemoryStore at cap; evicted weakest entry (F-MEM-3)"
218 );
219 }
220 }
221 let entry_tokens = build_entry_tokens(&entry);
224 let id = entry.id.clone();
225 if let Some(old_tokens) = tokens.get(&id) {
228 deindex_entry(&mut inverted, &id, old_tokens);
229 }
230 index_entry(&mut inverted, &id, &entry_tokens);
231 entries.insert(id.clone(), entry);
232 tokens.insert(id, entry_tokens);
233 Ok(())
234 })
235 }
236
237 fn recall(
238 &self,
239 scope: &TenantScope,
240 query: MemoryQuery,
241 ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryEntry>, Error>> + Send + '_>> {
242 let tenant_id = scope.tenant_id.clone();
243 Box::pin(async move {
244 let mut entries = self.entries.write();
249 let tokens_cache = self.tokens.read();
254 let inverted = self.inverted.read();
258
259 let now = Utc::now();
260 let query_tokens: Vec<String> = query
261 .text
262 .as_deref()
263 .map(|t| {
264 let mut seen = std::collections::HashSet::new();
265 t.to_lowercase()
266 .split_whitespace()
267 .filter(|tok| seen.insert(tok.to_string()))
268 .map(String::from)
269 .collect()
270 })
271 .unwrap_or_default();
272
273 let exact_word_candidate_ids: Option<Vec<String>> =
288 if query.exact_words && !query_tokens.is_empty() {
289 let mut ids: HashSet<String> = HashSet::new();
290 for token in &query_tokens {
291 if let Some(bucket) = inverted.get(token.as_str()) {
292 ids.extend(bucket.iter().cloned());
293 }
294 }
295 Some(ids.into_iter().collect())
296 } else {
297 None
298 };
299
300 let filter_logic = |e: &MemoryEntry| -> bool {
312 if e.author_tenant_id.as_deref() != Some(tenant_id.as_str()) {
314 return false;
315 }
316 if let Some(ref agent) = query.agent {
317 if e.agent != *agent {
318 return false;
319 }
320 } else if let Some(ref prefix) = query.agent_prefix
321 && !e.agent.starts_with(prefix.as_str())
322 {
323 return false;
324 }
325 if let Some(ref cat) = query.category
326 && e.category != *cat
327 {
328 return false;
329 }
330 if !query.tags.is_empty() && !query.tags.iter().any(|t| e.tags.contains(t)) {
331 return false;
332 }
333 if let Some(ref mt) = query.memory_type
334 && e.memory_type != *mt
335 {
336 return false;
337 }
338 if let Some(max_conf) = query.max_confidentiality
339 && e.confidentiality > max_conf
340 {
341 return false;
342 }
343 if let Some(min_s) = query.min_strength {
344 let eff =
345 effective_strength(e.strength, e.last_accessed, now, STRENGTH_DECAY_RATE);
346 if eff < min_s {
347 return false;
348 }
349 }
350 true
351 };
352
353 let candidates: Vec<(&MemoryEntry, &EntryTokens)> = match exact_word_candidate_ids {
354 Some(ids) => ids
355 .iter()
356 .filter_map(|id| {
357 let e = entries.get(id)?;
358 if !filter_logic(e) {
359 return None;
360 }
361 let tokens = tokens_cache.get(id)?;
362 Some((e, tokens))
363 })
364 .collect(),
365 None => entries
366 .values()
367 .filter_map(|e| {
368 if !filter_logic(e) {
369 return None;
370 }
371 let tokens = tokens_cache.get(&e.id)?;
376 if !query_tokens.is_empty() {
379 let has_match = query_tokens.iter().any(|token| {
380 tokens.lower_content.contains(token.as_str())
381 || tokens
382 .lower_keywords
383 .iter()
384 .any(|k| k.contains(token.as_str()))
385 });
386 if !has_match {
387 return None;
388 }
389 }
390 Some((e, tokens))
391 })
392 .collect(),
393 };
394
395 let avgdl = if candidates.is_empty() {
398 1.0
399 } else {
400 let total_words: usize =
401 candidates.iter().map(|(_, t)| t.content_words.len()).sum();
402 (total_words as f64 / candidates.len() as f64).max(1.0)
403 };
404
405 let bm25_map: HashMap<&str, f64> = candidates
409 .iter()
410 .map(|(e, t)| {
411 let score = bm25::bm25_score_pre(
412 &t.content_words,
413 &t.lower_keywords,
414 &query_tokens,
415 avgdl,
416 bm25::DEFAULT_K1,
417 bm25::DEFAULT_B,
418 );
419 (e.id.as_str(), score)
420 })
421 .collect();
422
423 let relevance_map: HashMap<&str, f64> = if let Some(ref q_emb) = query.query_embedding {
426 let mut bm25_ranked: Vec<(&str, f64)> =
428 bm25_map.iter().map(|(id, &s)| (*id, s)).collect();
429 bm25_ranked
430 .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
431
432 let mut vector_ranked: Vec<(&str, f64)> = candidates
434 .iter()
435 .filter_map(|(e, _)| {
436 e.embedding
437 .as_ref()
438 .map(|emb| (e.id.as_str(), hybrid::cosine_similarity(emb, q_emb)))
439 })
440 .collect();
441 vector_ranked
442 .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
443
444 if vector_ranked.is_empty() {
445 let max_bm25 = bm25_map
447 .values()
448 .copied()
449 .fold(f64::NEG_INFINITY, f64::max)
450 .max(1.0);
451 bm25_map
452 .iter()
453 .map(|(id, &s)| (*id, s / max_bm25))
454 .collect()
455 } else {
456 let fused = hybrid::rrf_fuse(&bm25_ranked, &vector_ranked, 50);
457 let max_fused = fused
458 .iter()
459 .map(|(_, s)| *s)
460 .fold(f64::NEG_INFINITY, f64::max)
461 .max(f64::EPSILON);
462 let mut out: HashMap<&str, f64> = HashMap::with_capacity(fused.len());
467 for (id_owned, score) in &fused {
468 if let Some((k, _)) = bm25_map.get_key_value(id_owned.as_str()) {
469 out.insert(k, score / max_fused);
470 }
471 }
472 out
473 }
474 } else {
475 let max_bm25 = bm25_map
477 .values()
478 .copied()
479 .fold(f64::NEG_INFINITY, f64::max)
480 .max(1.0);
481 bm25_map
482 .iter()
483 .map(|(id, &s)| (*id, s / max_bm25))
484 .collect()
485 };
486
487 let mut scored: Vec<(&MemoryEntry, f64)> = candidates
492 .iter()
493 .map(|(e, _)| {
494 let relevance = relevance_map.get(e.id.as_str()).copied().unwrap_or(0.0);
495 let eff =
496 effective_strength(e.strength, e.last_accessed, now, STRENGTH_DECAY_RATE);
497 let score = composite_score(
498 &self.scoring_weights,
499 e.created_at,
500 now,
501 e.importance,
502 relevance,
503 eff,
504 );
505 (*e, score)
506 })
507 .collect();
508 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
509
510 if query.limit > 0 && scored.len() > query.limit {
511 scored.truncate(query.limit);
512 }
513
514 let top_ids: std::collections::HashSet<String> =
517 scored.iter().map(|(e, _)| e.id.clone()).collect();
518 let mut to_expand = Vec::new();
519 let mut seen_expanded = std::collections::HashSet::new();
520 for (entry, _) in &scored {
521 for related_id in &entry.related_ids {
522 if !top_ids.contains(related_id) && seen_expanded.insert(related_id.clone()) {
523 to_expand.push(related_id.clone());
524 }
525 }
526 }
527
528 let min_s = query.min_strength.unwrap_or(0.0);
533 let max_bm25 = bm25_map
535 .values()
536 .copied()
537 .fold(f64::NEG_INFINITY, f64::max)
538 .max(1.0);
539 let mut expanded_added = 0usize;
540 for related_id in &to_expand {
541 if let Some(related) = entries.get(related_id) {
542 if let Some(max_conf) = query.max_confidentiality
543 && related.confidentiality > max_conf
544 {
545 continue;
546 }
547 let eff = effective_strength(
548 related.strength,
549 related.last_accessed,
550 now,
551 STRENGTH_DECAY_RATE,
552 );
553 if eff < min_s {
554 continue;
555 }
556 let Some(related_tokens) = tokens_cache.get(related_id) else {
561 continue;
562 };
563 let relevance = bm25::bm25_score_pre(
564 &related_tokens.content_words,
565 &related_tokens.lower_keywords,
566 &query_tokens,
567 avgdl,
568 bm25::DEFAULT_K1,
569 bm25::DEFAULT_B,
570 );
571 let normalised = relevance / max_bm25;
572 let score = composite_score(
573 &self.scoring_weights,
574 related.created_at,
575 now,
576 related.importance,
577 normalised,
578 eff,
579 );
580 scored.push((related, score));
581 expanded_added += 1;
582 }
583 }
584
585 if expanded_added > 0 {
586 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
587 if query.limit > 0 && scored.len() > query.limit {
588 scored.truncate(query.limit);
589 }
590 }
591
592 let top_ids_final: Vec<String> = scored.iter().map(|(e, _)| e.id.clone()).collect();
597 drop(scored);
598 drop(candidates);
599 drop(tokens_cache);
600 drop(inverted);
601
602 let reinforce = query.reinforce;
607 let mut results: Vec<MemoryEntry> = Vec::with_capacity(top_ids_final.len());
608 for id in top_ids_final {
609 if let Some(e) = entries.get_mut(&id) {
610 e.access_count += 1;
611 e.last_accessed = now;
612 if reinforce {
613 e.strength = (e.strength + 0.2).min(1.0);
615 }
616 results.push(e.clone());
617 }
618 }
619
620 Ok(results)
621 })
622 }
623
624 fn update(
625 &self,
626 scope: &TenantScope,
627 id: &str,
628 content: String,
629 ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
630 let id = id.to_string();
631 let tenant_id = scope.tenant_id.clone();
632 Box::pin(async move {
633 let mut entries = self.entries.write();
634 let mut tokens = self.tokens.write();
635 let mut inverted = self.inverted.write();
636 match entries.get_mut(&id) {
637 Some(entry) if entry.author_tenant_id.as_deref() == Some(tenant_id.as_str()) => {
638 entry.content = content;
639 entry.last_accessed = Utc::now();
640 let new_tokens = build_entry_tokens(entry);
643 if let Some(old_tokens) = tokens.get(&id) {
644 deindex_entry(&mut inverted, &id, old_tokens);
645 }
646 index_entry(&mut inverted, &id, &new_tokens);
647 tokens.insert(id.clone(), new_tokens);
648 Ok(())
649 }
650 Some(_) => {
651 Err(Error::Memory(format!("memory not found: {id}")))
653 }
654 None => Err(Error::Memory(format!("memory not found: {id}"))),
655 }
656 })
657 }
658
659 fn forget(
660 &self,
661 scope: &TenantScope,
662 id: &str,
663 ) -> Pin<Box<dyn Future<Output = Result<bool, Error>> + Send + '_>> {
664 let id = id.to_string();
665 let tenant_id = scope.tenant_id.clone();
666 Box::pin(async move {
667 let mut entries = self.entries.write();
668 let mut tokens = self.tokens.write();
669 let mut inverted = self.inverted.write();
670 let belongs = entries
674 .get(&id)
675 .map(|e| e.author_tenant_id.as_deref() == Some(tenant_id.as_str()))
676 .unwrap_or(false);
677 if belongs {
678 let removed = entries.remove(&id).is_some();
679 if let Some(old_tokens) = tokens.remove(&id) {
680 deindex_entry(&mut inverted, &id, &old_tokens);
681 }
682 Ok(removed)
683 } else {
684 Ok(false)
685 }
686 })
687 }
688
689 fn add_link(
690 &self,
691 scope: &TenantScope,
692 id: &str,
693 related_id: &str,
694 ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
695 let id = id.to_string();
696 let related_id = related_id.to_string();
697 let tenant_id = scope.tenant_id.clone();
698 Box::pin(async move {
699 let mut entries = self.entries.write();
700
701 let id_ok = entries
703 .get(&id)
704 .map(|e| e.author_tenant_id.as_deref() == Some(tenant_id.as_str()))
705 .unwrap_or(false);
706 let rel_ok = entries
707 .get(&related_id)
708 .map(|e| e.author_tenant_id.as_deref() == Some(tenant_id.as_str()))
709 .unwrap_or(false);
710
711 if id_ok
712 && let Some(entry) = entries.get_mut(&id)
713 && !entry.related_ids.contains(&related_id)
714 {
715 entry.related_ids.push(related_id.clone());
716 }
717 if rel_ok
718 && let Some(entry) = entries.get_mut(&related_id)
719 && !entry.related_ids.contains(&id)
720 {
721 entry.related_ids.push(id);
722 }
723 Ok(())
724 })
725 }
726
727 fn prune(
728 &self,
729 scope: &TenantScope,
730 min_strength: f64,
731 min_age: chrono::Duration,
732 agent_prefix: Option<&str>,
733 ) -> Pin<Box<dyn Future<Output = Result<usize, Error>> + Send + '_>> {
734 let owned_prefix = agent_prefix.map(String::from);
735 let tenant_id = scope.tenant_id.clone();
736 Box::pin(async move {
737 let mut entries = self.entries.write();
738
739 let now = Utc::now();
740 let to_remove: Vec<String> = entries
741 .values()
742 .filter(|e| {
743 if e.author_tenant_id.as_deref() != Some(tenant_id.as_str()) {
745 return false;
746 }
747 if let Some(ref prefix) = owned_prefix {
755 let p = prefix.as_str();
756 let agent = e.agent.as_str();
757 let separator_match = agent.len() > p.len()
758 && agent.starts_with(p)
759 && agent.as_bytes()[p.len()] == b':';
760 if agent != p && !separator_match {
761 return false;
762 }
763 }
764 let eff =
765 effective_strength(e.strength, e.last_accessed, now, STRENGTH_DECAY_RATE);
766 eff < min_strength && now.signed_duration_since(e.created_at) > min_age
767 })
768 .map(|e| e.id.clone())
769 .collect();
770
771 let count = to_remove.len();
772 let mut tokens = self.tokens.write();
773 let mut inverted = self.inverted.write();
774 for id in to_remove {
775 entries.remove(&id);
776 if let Some(old_tokens) = tokens.remove(&id) {
777 deindex_entry(&mut inverted, &id, &old_tokens);
778 }
779 }
780 Ok(count)
781 })
782 }
783}
784
785#[cfg(test)]
786mod tests {
787 use super::*;
788 use chrono::Utc;
789
790 use super::super::{Confidentiality, MemoryType};
791
792 fn test_scope() -> TenantScope {
793 TenantScope::default()
794 }
795
796 fn make_entry(id: &str, agent: &str, content: &str, category: &str) -> MemoryEntry {
797 MemoryEntry {
798 id: id.into(),
799 agent: agent.into(),
800 content: content.into(),
801 category: category.into(),
802 tags: vec![],
803 created_at: Utc::now(),
804 last_accessed: Utc::now(),
805 access_count: 0,
806 importance: 5,
807 memory_type: MemoryType::default(),
808 keywords: vec![],
809 summary: None,
810 strength: 1.0,
811 related_ids: vec![],
812 source_ids: vec![],
813 embedding: None,
814 confidentiality: Confidentiality::default(),
815 author_user_id: None,
816 author_tenant_id: None,
817 }
818 }
819
820 fn make_entry_with_tags(
821 id: &str,
822 agent: &str,
823 content: &str,
824 category: &str,
825 tags: Vec<String>,
826 ) -> MemoryEntry {
827 MemoryEntry {
828 id: id.into(),
829 agent: agent.into(),
830 content: content.into(),
831 category: category.into(),
832 tags,
833 created_at: Utc::now(),
834 last_accessed: Utc::now(),
835 access_count: 0,
836 importance: 5,
837 memory_type: MemoryType::default(),
838 keywords: vec![],
839 summary: None,
840 strength: 1.0,
841 related_ids: vec![],
842 source_ids: vec![],
843 embedding: None,
844 confidentiality: Confidentiality::default(),
845 author_user_id: None,
846 author_tenant_id: None,
847 }
848 }
849
850 #[tokio::test]
851 async fn store_and_recall() {
852 let store = InMemoryStore::new();
853 let entry = make_entry("m1", "agent1", "Rust is fast", "fact");
854 store.store(&test_scope(), entry).await.unwrap();
855
856 let results = store
857 .recall(
858 &test_scope(),
859 MemoryQuery {
860 limit: 10,
861 ..Default::default()
862 },
863 )
864 .await
865 .unwrap();
866 assert_eq!(results.len(), 1);
867 assert_eq!(results[0].content, "Rust is fast");
868 }
869
870 #[tokio::test]
871 async fn recall_by_text() {
872 let store = InMemoryStore::new();
873 store
874 .store(&test_scope(), make_entry("m1", "a", "Rust is fast", "fact"))
875 .await
876 .unwrap();
877 store
878 .store(
879 &test_scope(),
880 make_entry("m2", "a", "Python is slow", "fact"),
881 )
882 .await
883 .unwrap();
884
885 let results = store
886 .recall(
887 &test_scope(),
888 MemoryQuery {
889 text: Some("rust".into()),
890 limit: 10,
891 ..Default::default()
892 },
893 )
894 .await
895 .unwrap();
896 assert_eq!(results.len(), 1);
897 assert_eq!(results[0].id, "m1");
898 }
899
900 #[tokio::test]
901 async fn recall_by_category() {
902 let store = InMemoryStore::new();
903 store
904 .store(
905 &test_scope(),
906 make_entry("m1", "a", "remember this", "fact"),
907 )
908 .await
909 .unwrap();
910 store
911 .store(
912 &test_scope(),
913 make_entry("m2", "a", "I saw something", "observation"),
914 )
915 .await
916 .unwrap();
917
918 let results = store
919 .recall(
920 &test_scope(),
921 MemoryQuery {
922 category: Some("observation".into()),
923 limit: 10,
924 ..Default::default()
925 },
926 )
927 .await
928 .unwrap();
929 assert_eq!(results.len(), 1);
930 assert_eq!(results[0].id, "m2");
931 }
932
933 #[tokio::test]
934 async fn recall_by_tags() {
935 let store = InMemoryStore::new();
936 store
937 .store(
938 &test_scope(),
939 make_entry_with_tags(
940 "m1",
941 "a",
942 "Rust memory safety",
943 "fact",
944 vec!["rust".into(), "safety".into()],
945 ),
946 )
947 .await
948 .unwrap();
949 store
950 .store(
951 &test_scope(),
952 make_entry_with_tags(
953 "m2",
954 "a",
955 "Go is garbage collected",
956 "fact",
957 vec!["go".into()],
958 ),
959 )
960 .await
961 .unwrap();
962
963 let results = store
964 .recall(
965 &test_scope(),
966 MemoryQuery {
967 tags: vec!["rust".into()],
968 limit: 10,
969 ..Default::default()
970 },
971 )
972 .await
973 .unwrap();
974 assert_eq!(results.len(), 1);
975 assert_eq!(results[0].id, "m1");
976 }
977
978 #[tokio::test]
979 async fn recall_by_agent() {
980 let store = InMemoryStore::new();
981 store
982 .store(
983 &test_scope(),
984 make_entry("m1", "researcher", "data point", "fact"),
985 )
986 .await
987 .unwrap();
988 store
989 .store(
990 &test_scope(),
991 make_entry("m2", "coder", "code snippet", "procedure"),
992 )
993 .await
994 .unwrap();
995
996 let results = store
997 .recall(
998 &test_scope(),
999 MemoryQuery {
1000 agent: Some("researcher".into()),
1001 limit: 10,
1002 ..Default::default()
1003 },
1004 )
1005 .await
1006 .unwrap();
1007 assert_eq!(results.len(), 1);
1008 assert_eq!(results[0].id, "m1");
1009 }
1010
1011 #[tokio::test]
1012 async fn recall_limit() {
1013 let store = InMemoryStore::new();
1014 for i in 0..10 {
1015 store
1016 .store(
1017 &test_scope(),
1018 make_entry(&format!("m{i}"), "a", &format!("entry {i}"), "fact"),
1019 )
1020 .await
1021 .unwrap();
1022 }
1023
1024 let results = store
1025 .recall(
1026 &test_scope(),
1027 MemoryQuery {
1028 limit: 3,
1029 ..Default::default()
1030 },
1031 )
1032 .await
1033 .unwrap();
1034 assert_eq!(results.len(), 3);
1035 }
1036
1037 #[tokio::test]
1038 async fn update_existing() {
1039 let store = InMemoryStore::new();
1040 store
1041 .store(&test_scope(), make_entry("m1", "a", "original", "fact"))
1042 .await
1043 .unwrap();
1044
1045 store
1046 .update(&test_scope(), "m1", "updated content".into())
1047 .await
1048 .unwrap();
1049
1050 let results = store
1051 .recall(
1052 &test_scope(),
1053 MemoryQuery {
1054 limit: 10,
1055 ..Default::default()
1056 },
1057 )
1058 .await
1059 .unwrap();
1060 assert_eq!(results[0].content, "updated content");
1061 }
1062
1063 #[tokio::test]
1064 async fn update_nonexistent() {
1065 let store = InMemoryStore::new();
1066 let err = store
1067 .update(&test_scope(), "missing", "content".into())
1068 .await
1069 .unwrap_err();
1070 assert!(err.to_string().contains("not found"));
1071 }
1072
1073 #[tokio::test]
1074 async fn forget_existing() {
1075 let store = InMemoryStore::new();
1076 store
1077 .store(&test_scope(), make_entry("m1", "a", "to delete", "fact"))
1078 .await
1079 .unwrap();
1080
1081 assert!(store.forget(&test_scope(), "m1").await.unwrap());
1082
1083 let results = store
1084 .recall(
1085 &test_scope(),
1086 MemoryQuery {
1087 limit: 10,
1088 ..Default::default()
1089 },
1090 )
1091 .await
1092 .unwrap();
1093 assert!(results.is_empty());
1094 }
1095
1096 #[tokio::test]
1097 async fn forget_nonexistent() {
1098 let store = InMemoryStore::new();
1099 assert!(!store.forget(&test_scope(), "missing").await.unwrap());
1100 }
1101
1102 #[test]
1103 fn is_send_sync() {
1104 fn assert_send_sync<T: Send + Sync>() {}
1105 assert_send_sync::<InMemoryStore>();
1106 }
1107
1108 #[tokio::test]
1109 async fn recall_sorts_by_composite_score() {
1110 let store = InMemoryStore::new();
1111
1112 let mut high_imp = make_entry("m1", "a", "high importance", "fact");
1114 high_imp.importance = 10;
1115 high_imp.created_at = Utc::now() - chrono::Duration::hours(48);
1116 store.store(&test_scope(), high_imp).await.unwrap();
1117
1118 let mut low_imp = make_entry("m2", "a", "low importance", "fact");
1120 low_imp.importance = 1;
1121 low_imp.created_at = Utc::now();
1122 store.store(&test_scope(), low_imp).await.unwrap();
1123
1124 let results = store
1125 .recall(
1126 &test_scope(),
1127 MemoryQuery {
1128 limit: 10,
1129 ..Default::default()
1130 },
1131 )
1132 .await
1133 .unwrap();
1134
1135 assert_eq!(results.len(), 2);
1136 assert_eq!(results[0].id, "m1");
1141 assert_eq!(results[1].id, "m2");
1142 }
1143
1144 #[tokio::test]
1145 async fn recall_recent_high_importance_first() {
1146 let store = InMemoryStore::new();
1147
1148 let mut old_low = make_entry("m1", "a", "old low", "fact");
1150 old_low.importance = 1;
1151 old_low.created_at = Utc::now() - chrono::Duration::hours(1000);
1152 store.store(&test_scope(), old_low).await.unwrap();
1153
1154 let mut recent_high = make_entry("m2", "a", "recent high", "fact");
1156 recent_high.importance = 10;
1157 recent_high.created_at = Utc::now();
1158 store.store(&test_scope(), recent_high).await.unwrap();
1159
1160 let results = store
1161 .recall(
1162 &test_scope(),
1163 MemoryQuery {
1164 limit: 10,
1165 ..Default::default()
1166 },
1167 )
1168 .await
1169 .unwrap();
1170
1171 assert_eq!(results[0].id, "m2");
1172 }
1173
1174 #[tokio::test]
1175 async fn recall_with_custom_weights() {
1176 let store = InMemoryStore::new().with_scoring_weights(ScoringWeights {
1178 alpha: 0.0,
1179 beta: 1.0,
1180 gamma: 0.0,
1181 delta: 0.0,
1182 decay_rate: 0.01,
1183 });
1184
1185 let mut low = make_entry("m1", "a", "recent but low", "fact");
1186 low.importance = 1;
1187 low.created_at = Utc::now();
1188 store.store(&test_scope(), low).await.unwrap();
1189
1190 let mut high = make_entry("m2", "a", "old but high", "fact");
1191 high.importance = 10;
1192 high.created_at = Utc::now() - chrono::Duration::hours(1000);
1193 store.store(&test_scope(), high).await.unwrap();
1194
1195 let results = store
1196 .recall(
1197 &test_scope(),
1198 MemoryQuery {
1199 limit: 10,
1200 ..Default::default()
1201 },
1202 )
1203 .await
1204 .unwrap();
1205
1206 assert_eq!(results[0].id, "m2");
1208 }
1209
1210 #[tokio::test]
1211 async fn recall_text_query_affects_relevance() {
1212 let store = InMemoryStore::new().with_scoring_weights(ScoringWeights {
1214 alpha: 0.0,
1215 beta: 0.0,
1216 gamma: 1.0,
1217 delta: 0.0,
1218 decay_rate: 0.01,
1219 });
1220
1221 let mut e1 = make_entry("m1", "a", "Rust is fast", "fact");
1222 e1.importance = 5;
1223 store.store(&test_scope(), e1).await.unwrap();
1224
1225 let results = store
1227 .recall(
1228 &test_scope(),
1229 MemoryQuery {
1230 text: Some("Rust".into()),
1231 limit: 10,
1232 ..Default::default()
1233 },
1234 )
1235 .await
1236 .unwrap();
1237
1238 assert_eq!(results.len(), 1);
1239 assert_eq!(results[0].id, "m1");
1240 }
1241
1242 #[tokio::test]
1243 async fn recall_limit_zero_returns_all() {
1244 let store = InMemoryStore::new();
1245 for i in 0..5 {
1246 store
1247 .store(
1248 &test_scope(),
1249 make_entry(&format!("m{i}"), "a", &format!("entry {i}"), "fact"),
1250 )
1251 .await
1252 .unwrap();
1253 }
1254
1255 let results = store
1257 .recall(
1258 &test_scope(),
1259 MemoryQuery {
1260 limit: 0,
1261 ..Default::default()
1262 },
1263 )
1264 .await
1265 .unwrap();
1266 assert_eq!(results.len(), 5);
1267 }
1268
1269 #[tokio::test]
1270 async fn recall_deduplicates_query_tokens() {
1271 let store = InMemoryStore::new().with_scoring_weights(ScoringWeights {
1274 alpha: 0.0,
1275 beta: 0.0,
1276 gamma: 1.0,
1277 delta: 0.0,
1278 decay_rate: 0.01,
1279 });
1280
1281 store
1282 .store(&test_scope(), make_entry("m1", "a", "Rust is fast", "fact"))
1283 .await
1284 .unwrap();
1285 store
1286 .store(
1287 &test_scope(),
1288 make_entry("m2", "a", "Python is slow", "fact"),
1289 )
1290 .await
1291 .unwrap();
1292
1293 let results = store
1295 .recall(
1296 &test_scope(),
1297 MemoryQuery {
1298 text: Some("rust rust rust".into()),
1299 limit: 10,
1300 ..Default::default()
1301 },
1302 )
1303 .await
1304 .unwrap();
1305
1306 assert_eq!(results.len(), 1);
1308 assert_eq!(results[0].id, "m1");
1309 }
1310
1311 #[tokio::test]
1312 async fn relevance_score_differentiates_results() {
1313 let store = InMemoryStore::new();
1317
1318 let mut entry_partial =
1319 make_entry("m1", "agent1", "Rust is popular in the industry", "fact");
1320 entry_partial.importance = 5;
1321
1322 let mut entry_full =
1323 make_entry("m2", "agent1", "Rust is fast and safe for systems", "fact");
1324 entry_full.importance = 5;
1325
1326 store.store(&test_scope(), entry_partial).await.unwrap();
1328 store.store(&test_scope(), entry_full).await.unwrap();
1329
1330 let results = store
1331 .recall(
1332 &test_scope(),
1333 MemoryQuery {
1334 text: Some("Rust fast".into()),
1335 limit: 10,
1336 ..Default::default()
1337 },
1338 )
1339 .await
1340 .unwrap();
1341
1342 assert_eq!(results.len(), 2);
1344 assert_eq!(
1346 results[0].id, "m2",
1347 "entry matching more query tokens should rank first"
1348 );
1349 }
1350
1351 #[tokio::test]
1354 async fn recall_filters_by_memory_type() {
1355 let store = InMemoryStore::new();
1356
1357 let mut episodic = make_entry("m1", "a", "episodic fact", "fact");
1358 episodic.memory_type = MemoryType::Episodic;
1359 store.store(&test_scope(), episodic).await.unwrap();
1360
1361 let mut semantic = make_entry("m2", "a", "semantic knowledge", "fact");
1362 semantic.memory_type = MemoryType::Semantic;
1363 store.store(&test_scope(), semantic).await.unwrap();
1364
1365 let mut reflection = make_entry("m3", "a", "reflection insight", "fact");
1366 reflection.memory_type = MemoryType::Reflection;
1367 store.store(&test_scope(), reflection).await.unwrap();
1368
1369 let results = store
1371 .recall(
1372 &test_scope(),
1373 MemoryQuery {
1374 memory_type: Some(MemoryType::Semantic),
1375 limit: 10,
1376 ..Default::default()
1377 },
1378 )
1379 .await
1380 .unwrap();
1381 assert_eq!(results.len(), 1);
1382 assert_eq!(results[0].id, "m2");
1383
1384 let results = store
1386 .recall(
1387 &test_scope(),
1388 MemoryQuery {
1389 memory_type: Some(MemoryType::Reflection),
1390 limit: 10,
1391 ..Default::default()
1392 },
1393 )
1394 .await
1395 .unwrap();
1396 assert_eq!(results.len(), 1);
1397 assert_eq!(results[0].id, "m3");
1398
1399 let results = store
1401 .recall(
1402 &test_scope(),
1403 MemoryQuery {
1404 limit: 10,
1405 ..Default::default()
1406 },
1407 )
1408 .await
1409 .unwrap();
1410 assert_eq!(results.len(), 3);
1411 }
1412
1413 #[tokio::test]
1414 async fn recall_filters_by_min_strength() {
1415 let store = InMemoryStore::new();
1416
1417 let mut strong = make_entry("m1", "a", "strong memory", "fact");
1418 strong.strength = 0.9;
1419 store.store(&test_scope(), strong).await.unwrap();
1420
1421 let mut weak = make_entry("m2", "a", "weak memory", "fact");
1422 weak.strength = 0.05;
1423 store.store(&test_scope(), weak).await.unwrap();
1424
1425 let results = store
1427 .recall(
1428 &test_scope(),
1429 MemoryQuery {
1430 min_strength: Some(0.5),
1431 limit: 10,
1432 ..Default::default()
1433 },
1434 )
1435 .await
1436 .unwrap();
1437 assert_eq!(results.len(), 1);
1438 assert_eq!(results[0].id, "m1");
1439 }
1440
1441 #[tokio::test]
1442 async fn strength_reinforced_on_access() {
1443 let store = InMemoryStore::new();
1444
1445 let mut entry = make_entry("m1", "a", "test", "fact");
1446 entry.strength = 0.5;
1447 store.store(&test_scope(), entry).await.unwrap();
1448
1449 let results = store
1451 .recall(
1452 &test_scope(),
1453 MemoryQuery {
1454 limit: 10,
1455 ..Default::default()
1456 },
1457 )
1458 .await
1459 .unwrap();
1460 assert!((results[0].strength - 0.7).abs() < f64::EPSILON);
1461
1462 let results = store
1464 .recall(
1465 &test_scope(),
1466 MemoryQuery {
1467 limit: 10,
1468 ..Default::default()
1469 },
1470 )
1471 .await
1472 .unwrap();
1473 assert!((results[0].strength - 0.9).abs() < f64::EPSILON);
1474 }
1475
1476 #[tokio::test]
1477 async fn store_preserves_caller_supplied_strength() {
1478 let store = InMemoryStore::new();
1482
1483 let mut weak = make_entry("m1", "a", "test", "fact");
1484 weak.strength = 0.05;
1485 store.store(&test_scope(), weak).await.unwrap();
1486
1487 let results = store
1488 .recall(
1489 &test_scope(),
1490 MemoryQuery {
1491 limit: 10,
1492 reinforce: false,
1493 ..Default::default()
1494 },
1495 )
1496 .await
1497 .unwrap();
1498 assert!(
1499 (results[0].strength - 0.05).abs() < f64::EPSILON,
1500 "store must preserve caller strength; recall(reinforce=false) must not mutate it (got {})",
1501 results[0].strength
1502 );
1503 }
1504
1505 #[tokio::test]
1506 async fn recall_with_reinforce_false_is_idempotent_for_strength() {
1507 let store = InMemoryStore::new();
1508
1509 let mut entry = make_entry("m1", "a", "test", "fact");
1510 entry.strength = 0.4;
1511 store.store(&test_scope(), entry).await.unwrap();
1512
1513 for _ in 0..5 {
1514 let results = store
1515 .recall(
1516 &test_scope(),
1517 MemoryQuery {
1518 limit: 10,
1519 reinforce: false,
1520 ..Default::default()
1521 },
1522 )
1523 .await
1524 .unwrap();
1525 assert!((results[0].strength - 0.4).abs() < f64::EPSILON);
1526 }
1527
1528 let results = store
1530 .recall(
1531 &test_scope(),
1532 MemoryQuery {
1533 limit: 10,
1534 reinforce: false,
1535 ..Default::default()
1536 },
1537 )
1538 .await
1539 .unwrap();
1540 assert!(results[0].access_count >= 5);
1541 }
1542
1543 #[tokio::test]
1544 async fn strength_capped_at_one() {
1545 let store = InMemoryStore::new();
1546
1547 let mut entry = make_entry("m1", "a", "test", "fact");
1548 entry.strength = 0.95;
1549 store.store(&test_scope(), entry).await.unwrap();
1550
1551 let results = store
1553 .recall(
1554 &test_scope(),
1555 MemoryQuery {
1556 limit: 10,
1557 ..Default::default()
1558 },
1559 )
1560 .await
1561 .unwrap();
1562 assert!((results[0].strength - 1.0).abs() < f64::EPSILON);
1563 }
1564
1565 #[tokio::test]
1566 async fn keywords_searched_during_recall() {
1567 let store = InMemoryStore::new();
1568
1569 let mut entry = make_entry("m1", "a", "Rust is great", "fact");
1571 entry.keywords = vec!["performance".into(), "speed".into()];
1572 store.store(&test_scope(), entry).await.unwrap();
1573
1574 let results = store
1575 .recall(
1576 &test_scope(),
1577 MemoryQuery {
1578 text: Some("performance".into()),
1579 limit: 10,
1580 ..Default::default()
1581 },
1582 )
1583 .await
1584 .unwrap();
1585 assert_eq!(results.len(), 1);
1586 assert_eq!(results[0].id, "m1");
1587 }
1588
1589 #[tokio::test]
1590 async fn add_link_bidirectional() {
1591 let store = InMemoryStore::new();
1592 store
1593 .store(&test_scope(), make_entry("m1", "a", "first", "fact"))
1594 .await
1595 .unwrap();
1596 store
1597 .store(&test_scope(), make_entry("m2", "a", "second", "fact"))
1598 .await
1599 .unwrap();
1600
1601 store.add_link(&test_scope(), "m1", "m2").await.unwrap();
1602
1603 let results = store
1604 .recall(
1605 &test_scope(),
1606 MemoryQuery {
1607 limit: 10,
1608 ..Default::default()
1609 },
1610 )
1611 .await
1612 .unwrap();
1613 let m1 = results.iter().find(|e| e.id == "m1").unwrap();
1614 let m2 = results.iter().find(|e| e.id == "m2").unwrap();
1615
1616 assert!(m1.related_ids.contains(&"m2".to_string()));
1617 assert!(m2.related_ids.contains(&"m1".to_string()));
1618 }
1619
1620 #[tokio::test]
1621 async fn add_link_idempotent() {
1622 let store = InMemoryStore::new();
1623 store
1624 .store(&test_scope(), make_entry("m1", "a", "first", "fact"))
1625 .await
1626 .unwrap();
1627 store
1628 .store(&test_scope(), make_entry("m2", "a", "second", "fact"))
1629 .await
1630 .unwrap();
1631
1632 store.add_link(&test_scope(), "m1", "m2").await.unwrap();
1634 store.add_link(&test_scope(), "m1", "m2").await.unwrap();
1635
1636 let results = store
1637 .recall(
1638 &test_scope(),
1639 MemoryQuery {
1640 limit: 10,
1641 ..Default::default()
1642 },
1643 )
1644 .await
1645 .unwrap();
1646 let m1 = results.iter().find(|e| e.id == "m1").unwrap();
1647 assert_eq!(
1648 m1.related_ids.iter().filter(|id| *id == "m2").count(),
1649 1,
1650 "should not have duplicate links"
1651 );
1652 }
1653
1654 #[tokio::test]
1655 async fn prune_removes_below_threshold() {
1656 let store = InMemoryStore::new();
1657
1658 let mut strong = make_entry("m1", "a", "strong", "fact");
1659 strong.strength = 0.8;
1660 strong.created_at = Utc::now() - chrono::Duration::hours(48);
1661 store.store(&test_scope(), strong).await.unwrap();
1662
1663 let mut weak = make_entry("m2", "a", "weak", "fact");
1664 weak.strength = 0.05;
1665 weak.created_at = Utc::now() - chrono::Duration::hours(48);
1666 store.store(&test_scope(), weak).await.unwrap();
1667
1668 let pruned = store
1669 .prune(&test_scope(), 0.1, chrono::Duration::hours(1), None)
1670 .await
1671 .unwrap();
1672 assert_eq!(pruned, 1);
1673
1674 let results = store
1675 .recall(
1676 &test_scope(),
1677 MemoryQuery {
1678 limit: 10,
1679 ..Default::default()
1680 },
1681 )
1682 .await
1683 .unwrap();
1684 assert_eq!(results.len(), 1);
1685 assert_eq!(results[0].id, "m1");
1686 }
1687
1688 #[tokio::test]
1689 async fn prune_respects_min_age() {
1690 let store = InMemoryStore::new();
1691
1692 let mut weak_recent = make_entry("m1", "a", "weak recent", "fact");
1694 weak_recent.strength = 0.01;
1695 weak_recent.created_at = Utc::now(); store.store(&test_scope(), weak_recent).await.unwrap();
1697
1698 let pruned = store
1699 .prune(&test_scope(), 0.1, chrono::Duration::hours(24), None)
1700 .await
1701 .unwrap();
1702 assert_eq!(pruned, 0, "recent entry should not be pruned");
1703 }
1704
1705 #[tokio::test]
1706 async fn prune_uses_effective_strength_with_decay() {
1707 let store = InMemoryStore::new();
1708
1709 let mut old_accessed = make_entry("m1", "a", "old accessed", "fact");
1712 old_accessed.strength = 0.5;
1713 old_accessed.created_at = Utc::now() - chrono::Duration::hours(30 * 24);
1714 old_accessed.last_accessed = Utc::now() - chrono::Duration::hours(30 * 24);
1715 store.store(&test_scope(), old_accessed).await.unwrap();
1716
1717 let mut recently_accessed = make_entry("m2", "a", "recently accessed", "fact");
1719 recently_accessed.strength = 0.5;
1720 recently_accessed.created_at = Utc::now() - chrono::Duration::hours(30 * 24);
1721 recently_accessed.last_accessed = Utc::now();
1722 store.store(&test_scope(), recently_accessed).await.unwrap();
1723
1724 let pruned = store
1728 .prune(&test_scope(), 0.1, chrono::Duration::hours(24), None)
1729 .await
1730 .unwrap();
1731 assert_eq!(pruned, 1, "old unaccessed entry should be pruned");
1732
1733 let results = store
1734 .recall(
1735 &test_scope(),
1736 MemoryQuery {
1737 limit: 10,
1738 ..Default::default()
1739 },
1740 )
1741 .await
1742 .unwrap();
1743 assert_eq!(results.len(), 1);
1744 assert_eq!(results[0].id, "m2");
1745 }
1746
1747 #[tokio::test]
1748 async fn prune_with_agent_prefix_only_removes_matching_agent() {
1749 let store = InMemoryStore::new();
1750
1751 let mut weak_a = make_entry("m1", "agent_a", "weak from A", "fact");
1753 weak_a.strength = 0.01;
1754 weak_a.created_at = Utc::now() - chrono::Duration::hours(48);
1755 weak_a.last_accessed = Utc::now() - chrono::Duration::hours(48);
1756 store.store(&test_scope(), weak_a).await.unwrap();
1757
1758 let mut weak_b = make_entry("m2", "agent_b", "weak from B", "fact");
1759 weak_b.strength = 0.01;
1760 weak_b.created_at = Utc::now() - chrono::Duration::hours(48);
1761 weak_b.last_accessed = Utc::now() - chrono::Duration::hours(48);
1762 store.store(&test_scope(), weak_b).await.unwrap();
1763
1764 let pruned = store
1766 .prune(
1767 &test_scope(),
1768 0.1,
1769 chrono::Duration::hours(1),
1770 Some("agent_a"),
1771 )
1772 .await
1773 .unwrap();
1774 assert_eq!(pruned, 1, "should only prune agent_a's entry");
1775
1776 let results = store
1777 .recall(
1778 &test_scope(),
1779 MemoryQuery {
1780 limit: 10,
1781 ..Default::default()
1782 },
1783 )
1784 .await
1785 .unwrap();
1786 assert_eq!(results.len(), 1);
1787 assert_eq!(results[0].id, "m2");
1788 assert_eq!(results[0].agent, "agent_b");
1789 }
1790
1791 #[tokio::test]
1797 async fn prune_does_not_match_overlapping_agent_prefix() {
1798 let store = InMemoryStore::new();
1799
1800 let mut weak_alice = make_entry("ma", "user:alice", "weak alice", "fact");
1802 weak_alice.strength = 0.01;
1803 weak_alice.created_at = Utc::now() - chrono::Duration::hours(48);
1804 weak_alice.last_accessed = Utc::now() - chrono::Duration::hours(48);
1805 store.store(&test_scope(), weak_alice).await.unwrap();
1806
1807 let mut weak_alice2 = make_entry("m2", "user:alice2", "weak alice2", "fact");
1809 weak_alice2.strength = 0.01;
1810 weak_alice2.created_at = Utc::now() - chrono::Duration::hours(48);
1811 weak_alice2.last_accessed = Utc::now() - chrono::Duration::hours(48);
1812 store.store(&test_scope(), weak_alice2).await.unwrap();
1813
1814 let mut weak_subagent = make_entry("ms", "user:alice:tool", "weak sub", "fact");
1816 weak_subagent.strength = 0.01;
1817 weak_subagent.created_at = Utc::now() - chrono::Duration::hours(48);
1818 weak_subagent.last_accessed = Utc::now() - chrono::Duration::hours(48);
1819 store.store(&test_scope(), weak_subagent).await.unwrap();
1820
1821 let pruned = store
1822 .prune(
1823 &test_scope(),
1824 0.1,
1825 chrono::Duration::hours(1),
1826 Some("user:alice"),
1827 )
1828 .await
1829 .unwrap();
1830
1831 assert_eq!(pruned, 2, "must prune alice and user:alice:tool only");
1833
1834 let results = store
1835 .recall(
1836 &test_scope(),
1837 MemoryQuery {
1838 limit: 10,
1839 ..Default::default()
1840 },
1841 )
1842 .await
1843 .unwrap();
1844 let agents: std::collections::HashSet<&str> =
1845 results.iter().map(|e| e.agent.as_str()).collect();
1846 assert!(
1847 agents.contains("user:alice2"),
1848 "alice2 must survive (overlapping prefix bypass): got {agents:?}"
1849 );
1850 }
1851
1852 #[tokio::test]
1853 async fn prune_none_prefix_removes_all_matching() {
1854 let store = InMemoryStore::new();
1855
1856 let mut weak_a = make_entry("m1", "agent_a", "weak from A", "fact");
1857 weak_a.strength = 0.01;
1858 weak_a.created_at = Utc::now() - chrono::Duration::hours(48);
1859 weak_a.last_accessed = Utc::now() - chrono::Duration::hours(48);
1860 store.store(&test_scope(), weak_a).await.unwrap();
1861
1862 let mut weak_b = make_entry("m2", "agent_b", "weak from B", "fact");
1863 weak_b.strength = 0.01;
1864 weak_b.created_at = Utc::now() - chrono::Duration::hours(48);
1865 weak_b.last_accessed = Utc::now() - chrono::Duration::hours(48);
1866 store.store(&test_scope(), weak_b).await.unwrap();
1867
1868 let pruned = store
1870 .prune(&test_scope(), 0.1, chrono::Duration::hours(1), None)
1871 .await
1872 .unwrap();
1873 assert_eq!(pruned, 2, "should prune all weak entries");
1874 }
1875
1876 #[tokio::test]
1877 async fn recall_bm25_ranks_better_than_naive_keyword() {
1878 let store = InMemoryStore::new().with_scoring_weights(ScoringWeights {
1881 alpha: 0.0,
1882 beta: 0.0,
1883 gamma: 1.0,
1884 delta: 0.0,
1885 decay_rate: 0.01,
1886 });
1887
1888 let e1 = make_entry("m1", "a", "Rust is a programming language", "fact");
1890 store.store(&test_scope(), e1).await.unwrap();
1891
1892 let e2 = make_entry(
1894 "m2",
1895 "a",
1896 "Rust has excellent performance and speed",
1897 "fact",
1898 );
1899 store.store(&test_scope(), e2).await.unwrap();
1900
1901 let results = store
1902 .recall(
1903 &test_scope(),
1904 MemoryQuery {
1905 text: Some("Rust performance".into()),
1906 limit: 10,
1907 ..Default::default()
1908 },
1909 )
1910 .await
1911 .unwrap();
1912
1913 assert_eq!(results.len(), 2);
1914 assert_eq!(
1915 results[0].id, "m2",
1916 "BM25 should rank entry matching more query terms first"
1917 );
1918 }
1919
1920 #[tokio::test]
1921 async fn recall_bm25_keyword_field_boosts_ranking() {
1922 let store = InMemoryStore::new().with_scoring_weights(ScoringWeights {
1924 alpha: 0.0,
1925 beta: 0.0,
1926 gamma: 1.0,
1927 delta: 0.0,
1928 decay_rate: 0.01,
1929 });
1930
1931 let e1 = make_entry("m1", "a", "optimization techniques for databases", "fact");
1933 store.store(&test_scope(), e1).await.unwrap();
1934
1935 let mut e2 = make_entry("m2", "a", "optimization techniques for systems", "fact");
1937 e2.keywords = vec!["optimization".into(), "databases".into()];
1938 store.store(&test_scope(), e2).await.unwrap();
1939
1940 let results = store
1941 .recall(
1942 &test_scope(),
1943 MemoryQuery {
1944 text: Some("optimization databases".into()),
1945 limit: 10,
1946 ..Default::default()
1947 },
1948 )
1949 .await
1950 .unwrap();
1951
1952 assert_eq!(results.len(), 2);
1953 assert_eq!(
1954 results[0].id, "m2",
1955 "entry with keyword match should rank higher"
1956 );
1957 }
1958
1959 #[tokio::test]
1960 async fn strength_affects_ranking() {
1961 let store = InMemoryStore::new().with_scoring_weights(ScoringWeights {
1963 alpha: 0.0,
1964 beta: 0.0,
1965 gamma: 0.0,
1966 delta: 1.0,
1967 decay_rate: 0.01,
1968 });
1969
1970 let mut weak = make_entry("m1", "a", "weak entry", "fact");
1971 weak.strength = 0.2;
1972 store.store(&test_scope(), weak).await.unwrap();
1973
1974 let mut strong = make_entry("m2", "a", "strong entry", "fact");
1975 strong.strength = 0.9;
1976 store.store(&test_scope(), strong).await.unwrap();
1977
1978 let results = store
1979 .recall(
1980 &test_scope(),
1981 MemoryQuery {
1982 limit: 10,
1983 ..Default::default()
1984 },
1985 )
1986 .await
1987 .unwrap();
1988
1989 assert_eq!(results.len(), 2);
1990 assert_eq!(
1991 results[0].id, "m2",
1992 "stronger entry should rank first when delta=1.0"
1993 );
1994 }
1995
1996 #[tokio::test]
1997 async fn hybrid_recall_cosine_boosts_semantic_match() {
1998 let store = InMemoryStore::new().with_scoring_weights(ScoringWeights {
2001 alpha: 0.0,
2002 beta: 0.0,
2003 gamma: 1.0,
2004 delta: 0.0,
2005 decay_rate: 0.01,
2006 });
2007
2008 let e1 = make_entry("m1", "a", "Rust is fast", "fact");
2010 store.store(&test_scope(), e1).await.unwrap();
2011
2012 let mut e2 = make_entry(
2014 "m2",
2015 "a",
2016 "Systems programming language with safety",
2017 "fact",
2018 );
2019 e2.embedding = Some(vec![0.9, 0.1, 0.0]);
2020 store.store(&test_scope(), e2).await.unwrap();
2021
2022 let results = store
2024 .recall(
2025 &test_scope(),
2026 MemoryQuery {
2027 text: Some("rust".into()),
2028 query_embedding: Some(vec![0.9, 0.1, 0.0]),
2029 limit: 10,
2030 ..Default::default()
2031 },
2032 )
2033 .await
2034 .unwrap();
2035
2036 assert_eq!(results.len(), 1);
2040 assert_eq!(results[0].id, "m1");
2041 }
2042
2043 #[tokio::test]
2044 async fn hybrid_recall_fuses_bm25_and_vector() {
2045 let store = InMemoryStore::new().with_scoring_weights(ScoringWeights {
2049 alpha: 0.0,
2050 beta: 0.0,
2051 gamma: 1.0,
2052 delta: 0.0,
2053 decay_rate: 0.01,
2054 });
2055
2056 let mut e1 = make_entry("m1", "a", "Rust is fast and fast", "fact");
2058 e1.embedding = Some(vec![0.0, 0.0, 1.0]); store.store(&test_scope(), e1).await.unwrap();
2060
2061 let mut e2 = make_entry("m2", "a", "Rust has zero-cost abstractions", "fact");
2063 e2.embedding = Some(vec![0.95, 0.05, 0.0]); store.store(&test_scope(), e2).await.unwrap();
2065
2066 let mut e3 = make_entry("m3", "a", "Rust is a programming language", "fact");
2068 e3.embedding = Some(vec![0.5, 0.5, 0.0]); store.store(&test_scope(), e3).await.unwrap();
2070
2071 let results = store
2075 .recall(
2076 &test_scope(),
2077 MemoryQuery {
2078 text: Some("rust fast".into()),
2079 query_embedding: Some(vec![0.95, 0.05, 0.0]),
2080 limit: 10,
2081 ..Default::default()
2082 },
2083 )
2084 .await
2085 .unwrap();
2086
2087 assert_eq!(results.len(), 3);
2088 assert_eq!(
2093 results[0].id, "m2",
2094 "entry with highest cosine similarity should rank first in hybrid mode"
2095 );
2096 }
2097
2098 #[tokio::test]
2099 async fn hybrid_recall_bm25_fallback_when_no_embeddings() {
2100 let store = InMemoryStore::new().with_scoring_weights(ScoringWeights {
2103 alpha: 0.0,
2104 beta: 0.0,
2105 gamma: 1.0,
2106 delta: 0.0,
2107 decay_rate: 0.01,
2108 });
2109
2110 let e1 = make_entry("m1", "a", "Rust programming language", "fact");
2111 store.store(&test_scope(), e1).await.unwrap();
2112
2113 let e2 = make_entry("m2", "a", "Rust performance and speed", "fact");
2114 store.store(&test_scope(), e2).await.unwrap();
2115
2116 let results = store
2117 .recall(
2118 &test_scope(),
2119 MemoryQuery {
2120 text: Some("Rust performance".into()),
2121 query_embedding: Some(vec![0.5, 0.5, 0.0]),
2122 limit: 10,
2123 ..Default::default()
2124 },
2125 )
2126 .await
2127 .unwrap();
2128
2129 assert_eq!(results.len(), 2);
2130 assert_eq!(results[0].id, "m2");
2132 }
2133
2134 #[tokio::test]
2135 async fn recall_follows_related_ids_one_hop() {
2136 let store = InMemoryStore::new();
2139
2140 let mut m1 = make_entry("m1", "a", "Rust is fast", "fact");
2141 m1.related_ids = vec!["m2".into()];
2142 store.store(&test_scope(), m1).await.unwrap();
2143
2144 let mut m2 = make_entry("m2", "a", "Memory safety guarantees", "fact");
2145 m2.related_ids = vec!["m1".into()];
2146 store.store(&test_scope(), m2).await.unwrap();
2147
2148 let results = store
2149 .recall(
2150 &test_scope(),
2151 MemoryQuery {
2152 text: Some("rust".into()),
2153 limit: 10,
2154 ..Default::default()
2155 },
2156 )
2157 .await
2158 .unwrap();
2159
2160 assert_eq!(results.len(), 2);
2162 let ids: Vec<&str> = results.iter().map(|e| e.id.as_str()).collect();
2163 assert!(ids.contains(&"m1"), "direct match should be in results");
2164 assert!(
2165 ids.contains(&"m2"),
2166 "linked entry should be surfaced via graph expansion"
2167 );
2168 }
2169
2170 #[tokio::test]
2171 async fn recall_graph_expansion_respects_strength_threshold() {
2172 let store = InMemoryStore::new();
2173
2174 let mut m1 = make_entry("m1", "a", "Rust is fast", "fact");
2175 m1.related_ids = vec!["m2".into()];
2176 store.store(&test_scope(), m1).await.unwrap();
2177
2178 let mut m2 = make_entry("m2", "a", "Weak linked memory", "fact");
2180 m2.related_ids = vec!["m1".into()];
2181 m2.strength = 0.01;
2182 m2.last_accessed = Utc::now() - chrono::Duration::hours(720); store.store(&test_scope(), m2).await.unwrap();
2184
2185 let results = store
2186 .recall(
2187 &test_scope(),
2188 MemoryQuery {
2189 text: Some("rust".into()),
2190 min_strength: Some(0.1),
2191 limit: 10,
2192 ..Default::default()
2193 },
2194 )
2195 .await
2196 .unwrap();
2197
2198 assert_eq!(results.len(), 1);
2200 assert_eq!(results[0].id, "m1");
2201 }
2202
2203 #[tokio::test]
2204 async fn recall_graph_expansion_does_not_duplicate() {
2205 let store = InMemoryStore::new();
2206
2207 let mut m1 = make_entry("m1", "a", "Rust is fast", "fact");
2209 m1.related_ids = vec!["m2".into()];
2210 store.store(&test_scope(), m1).await.unwrap();
2211
2212 let mut m2 = make_entry("m2", "a", "Rust is safe", "fact");
2213 m2.related_ids = vec!["m1".into()];
2214 store.store(&test_scope(), m2).await.unwrap();
2215
2216 let results = store
2217 .recall(
2218 &test_scope(),
2219 MemoryQuery {
2220 text: Some("rust".into()),
2221 limit: 10,
2222 ..Default::default()
2223 },
2224 )
2225 .await
2226 .unwrap();
2227
2228 assert_eq!(results.len(), 2);
2230 let ids: Vec<&str> = results.iter().map(|e| e.id.as_str()).collect();
2231 assert_eq!(
2232 ids.iter().filter(|&&id| id == "m1").count(),
2233 1,
2234 "m1 should appear exactly once"
2235 );
2236 assert_eq!(
2237 ids.iter().filter(|&&id| id == "m2").count(),
2238 1,
2239 "m2 should appear exactly once"
2240 );
2241 }
2242
2243 #[tokio::test]
2244 async fn recall_agent_prefix_matches_sub_namespaces() {
2245 let store = InMemoryStore::new();
2246 store
2248 .store(
2249 &test_scope(),
2250 make_entry("m1", "tg:123:assistant", "likes Rust", "fact"),
2251 )
2252 .await
2253 .unwrap();
2254 store
2255 .store(
2256 &test_scope(),
2257 make_entry("m2", "tg:123:researcher", "loves coffee", "fact"),
2258 )
2259 .await
2260 .unwrap();
2261 store
2263 .store(
2264 &test_scope(),
2265 make_entry("m3", "tg:456:assistant", "prefers Python", "fact"),
2266 )
2267 .await
2268 .unwrap();
2269
2270 let results = store
2271 .recall(
2272 &test_scope(),
2273 MemoryQuery {
2274 agent_prefix: Some("tg:123".into()),
2275 limit: 10,
2276 ..Default::default()
2277 },
2278 )
2279 .await
2280 .unwrap();
2281
2282 assert_eq!(results.len(), 2);
2283 let ids: Vec<&str> = results.iter().map(|e| e.id.as_str()).collect();
2284 assert!(ids.contains(&"m1"));
2285 assert!(ids.contains(&"m2"));
2286 }
2287
2288 #[tokio::test]
2289 async fn recall_agent_exact_takes_precedence_over_prefix() {
2290 let store = InMemoryStore::new();
2291 store
2292 .store(
2293 &test_scope(),
2294 make_entry("m1", "tg:123:assistant", "from assistant", "fact"),
2295 )
2296 .await
2297 .unwrap();
2298 store
2299 .store(
2300 &test_scope(),
2301 make_entry("m2", "tg:123:researcher", "from researcher", "fact"),
2302 )
2303 .await
2304 .unwrap();
2305
2306 let results = store
2308 .recall(
2309 &test_scope(),
2310 MemoryQuery {
2311 agent: Some("tg:123:assistant".into()),
2312 agent_prefix: Some("tg:123".into()), limit: 10,
2314 ..Default::default()
2315 },
2316 )
2317 .await
2318 .unwrap();
2319
2320 assert_eq!(results.len(), 1);
2321 assert_eq!(results[0].id, "m1");
2322 }
2323
2324 #[tokio::test]
2325 async fn recall_filters_by_max_confidentiality() {
2326 use super::super::Confidentiality;
2327 let store = InMemoryStore::new();
2328
2329 let mut public = make_entry("m1", "a", "public fact", "fact");
2330 public.confidentiality = Confidentiality::Public;
2331 store.store(&test_scope(), public).await.unwrap();
2332
2333 let mut internal = make_entry("m2", "a", "internal note", "fact");
2334 internal.confidentiality = Confidentiality::Internal;
2335 store.store(&test_scope(), internal).await.unwrap();
2336
2337 let mut confidential = make_entry("m3", "a", "private expense", "fact");
2338 confidential.confidentiality = Confidentiality::Confidential;
2339 store.store(&test_scope(), confidential).await.unwrap();
2340
2341 let mut restricted = make_entry("m4", "a", "api key", "fact");
2342 restricted.confidentiality = Confidentiality::Restricted;
2343 store.store(&test_scope(), restricted).await.unwrap();
2344
2345 let results = store
2347 .recall(
2348 &test_scope(),
2349 MemoryQuery {
2350 max_confidentiality: Some(Confidentiality::Public),
2351 ..Default::default()
2352 },
2353 )
2354 .await
2355 .unwrap();
2356 assert_eq!(results.len(), 1);
2357 assert_eq!(results[0].id, "m1");
2358
2359 let results = store
2361 .recall(
2362 &test_scope(),
2363 MemoryQuery {
2364 max_confidentiality: None,
2365 ..Default::default()
2366 },
2367 )
2368 .await
2369 .unwrap();
2370 assert_eq!(results.len(), 4);
2371
2372 let results = store
2374 .recall(
2375 &test_scope(),
2376 MemoryQuery {
2377 max_confidentiality: Some(Confidentiality::Confidential),
2378 ..Default::default()
2379 },
2380 )
2381 .await
2382 .unwrap();
2383 assert_eq!(results.len(), 3);
2384 assert!(results.iter().all(|e| e.id != "m4"));
2385 }
2386
2387 #[tokio::test]
2388 async fn graph_expansion_respects_max_confidentiality() {
2389 use super::super::Confidentiality;
2390 let store = InMemoryStore::new();
2391
2392 let mut public = make_entry("m1", "a", "project update", "fact");
2394 public.confidentiality = Confidentiality::Public;
2395 public.related_ids = vec!["m2".into()];
2396 public.keywords = vec!["project".into()];
2397 store.store(&test_scope(), public).await.unwrap();
2398
2399 let mut confidential = make_entry("m2", "a", "private expense data", "fact");
2400 confidential.confidentiality = Confidentiality::Confidential;
2401 confidential.keywords = vec!["expense".into()];
2402 store.store(&test_scope(), confidential).await.unwrap();
2403
2404 let results = store
2406 .recall(
2407 &test_scope(),
2408 MemoryQuery {
2409 text: Some("project".into()),
2410 max_confidentiality: Some(Confidentiality::Public),
2411 ..Default::default()
2412 },
2413 )
2414 .await
2415 .unwrap();
2416
2417 assert!(
2418 results.iter().all(|e| e.id != "m2"),
2419 "graph expansion should not include Confidential entries when capped at Public"
2420 );
2421 assert!(results.iter().any(|e| e.id == "m1"));
2422 }
2423
2424 #[tokio::test]
2427 async fn recall_does_not_leak_across_tenants() {
2428 let store = InMemoryStore::new();
2429 let acme = TenantScope::new("acme");
2430 let globex = TenantScope::new("globex");
2431
2432 store
2433 .store(&acme, make_entry("a1", "agent", "acme-secret", "fact"))
2434 .await
2435 .unwrap();
2436 store
2437 .store(&globex, make_entry("g1", "agent", "globex-secret", "fact"))
2438 .await
2439 .unwrap();
2440
2441 let acme_results = store
2442 .recall(
2443 &acme,
2444 MemoryQuery {
2445 agent: Some("agent".into()),
2446 ..Default::default()
2447 },
2448 )
2449 .await
2450 .unwrap();
2451 assert_eq!(acme_results.len(), 1);
2452 assert_eq!(acme_results[0].id, "a1");
2453
2454 let globex_results = store
2455 .recall(
2456 &globex,
2457 MemoryQuery {
2458 agent: Some("agent".into()),
2459 ..Default::default()
2460 },
2461 )
2462 .await
2463 .unwrap();
2464 assert_eq!(globex_results.len(), 1);
2465 assert_eq!(globex_results[0].id, "g1");
2466 }
2467
2468 #[tokio::test]
2469 async fn forget_does_not_delete_other_tenant() {
2470 let store = InMemoryStore::new();
2471 let acme = TenantScope::new("acme");
2472 let globex = TenantScope::new("globex");
2473
2474 store
2475 .store(&acme, make_entry("a1", "agent", "x", "fact"))
2476 .await
2477 .unwrap();
2478 store
2479 .store(&globex, make_entry("g1", "agent", "y", "fact"))
2480 .await
2481 .unwrap();
2482
2483 let removed = store.forget(&globex, "a1").await.unwrap();
2485 assert!(!removed);
2486
2487 let acme_results = store
2488 .recall(
2489 &acme,
2490 MemoryQuery {
2491 agent: Some("agent".into()),
2492 ..Default::default()
2493 },
2494 )
2495 .await
2496 .unwrap();
2497 assert_eq!(acme_results.len(), 1);
2498 }
2499
2500 #[tokio::test]
2501 async fn update_does_not_modify_other_tenant() {
2502 let store = InMemoryStore::new();
2503 let acme = TenantScope::new("acme");
2504 let globex = TenantScope::new("globex");
2505
2506 store
2507 .store(&acme, make_entry("a1", "agent", "original", "fact"))
2508 .await
2509 .unwrap();
2510
2511 let err = store
2513 .update(&globex, "a1", "tampered".into())
2514 .await
2515 .unwrap_err();
2516 assert!(err.to_string().contains("memory not found"), "got: {err}");
2517
2518 let results = store
2520 .recall(
2521 &acme,
2522 MemoryQuery {
2523 agent: Some("agent".into()),
2524 ..Default::default()
2525 },
2526 )
2527 .await
2528 .unwrap();
2529 assert_eq!(results.len(), 1);
2530 assert_eq!(results[0].content, "original");
2531 }
2532
2533 #[tokio::test]
2534 async fn store_under_scope_populates_author_tenant_id() {
2535 let store = InMemoryStore::new();
2536 let scope = TenantScope::new("acme").with_user("u-42");
2537 store
2538 .store(&scope, make_entry("s1", "agent", "x", "fact"))
2539 .await
2540 .unwrap();
2541
2542 let results = store
2543 .recall(
2544 &scope,
2545 MemoryQuery {
2546 agent: Some("agent".into()),
2547 ..Default::default()
2548 },
2549 )
2550 .await
2551 .unwrap();
2552 assert_eq!(results.len(), 1);
2553 assert_eq!(results[0].author_tenant_id.as_deref(), Some("acme"));
2554 assert_eq!(results[0].author_user_id.as_deref(), Some("u-42"));
2555 }
2556}