1use lmdb::{Cursor, Database, Environment, Transaction, WriteFlags};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::sync::{Arc, RwLock};
7use wm_core::{
8 CoreError, EpisodicCapturePolicy, EpisodicId, EpisodicKind, EpisodicRecord, MemoryTransition,
9 Result, ValidityState,
10};
11
12use crate::embedder::Embedder;
13use crate::enrichment::VocabularyEnrichment;
14use crate::episodic_keys::{AdaptiveAliases, key_index_terms_with_aliases};
15use crate::query_planner::QueryPlan;
16use crate::search::strip_stopwords;
17
18#[cfg(test)]
25#[derive(Clone, Copy, PartialEq, Eq)]
26enum CommitBoundary {
27 BeforeRawCommit,
28 AfterRawCommit,
29}
30
31#[cfg(test)]
32fn commit_boundary_test_hook(boundary: CommitBoundary, records: &[EpisodicRecord]) -> Result<()> {
33 use std::io::{Read, Write};
34 let Ok(case) = std::env::var("WM_Q06_CASE") else {
35 return Ok(());
36 };
37 let expected = match boundary {
38 CommitBoundary::BeforeRawCommit => "before_raw_commit",
39 CommitBoundary::AfterRawCommit => "after_raw_commit",
40 };
41 if case != expected {
42 return Ok(());
43 }
44 let uuid = std::env::var("WM_Q06_UUID")
45 .map_err(|_| CoreError::Memory("q06 hook: WM_Q06_UUID is required".into()))?;
46 let names: Vec<String> = records.iter().map(|record| record.id.to_string()).collect();
47 if names.len() != 1 || names[0] != uuid {
48 return Err(CoreError::Memory(
49 "q06 hook: candidate set must be the single WM_Q06_UUID record".into(),
50 ));
51 }
52 println!("WM_Q06_BOUNDARY {expected} {uuid}");
53 let _ = std::io::stdout().flush();
54 let mut release = [0_u8; 1];
55 match std::io::stdin().read(&mut release) {
56 Ok(0) => Err(CoreError::Memory(
57 "q06 hook: stdin closed before release".into(),
58 )),
59 Ok(_) => Ok(()),
60 Err(e) => Err(CoreError::Memory(format!(
61 "q06 hook: stdin read failed: {e}"
62 ))),
63 }
64}
65
66#[derive(Debug, Clone)]
68pub struct EpisodicSearchResult {
69 pub record: EpisodicRecord,
70 pub score: f32,
71 pub matched_terms: usize,
72}
73
74pub struct EpisodicStore<'a> {
76 env: &'a Environment,
77 db: Database,
78 term_db: Database,
79 term_cache: Arc<RwLock<HashMap<String, Vec<EpisodicId>>>>,
80 mutation_count: &'a std::sync::atomic::AtomicU64,
81 embedder: Option<Arc<dyn Embedder + Send + Sync>>,
82 aliases: Option<AdaptiveAliases>,
83 enrichment: Option<VocabularyEnrichment>,
84}
85
86impl<'a> EpisodicStore<'a> {
87 pub(crate) fn new(
88 env: &'a Environment,
89 db: Database,
90 term_db: Database,
91 term_cache: Arc<RwLock<HashMap<String, Vec<EpisodicId>>>>,
92 mutation_count: &'a std::sync::atomic::AtomicU64,
93 ) -> Self {
94 Self {
95 env,
96 db,
97 term_db,
98 term_cache,
99 mutation_count,
100 embedder: None,
101 aliases: None,
102 enrichment: None,
103 }
104 }
105
106 #[must_use]
108 pub fn with_adaptive_aliases(mut self, aliases: AdaptiveAliases) -> Self {
109 if !aliases.is_empty() {
110 self.aliases = Some(aliases);
111 }
112 self
113 }
114
115 #[must_use]
117 pub fn with_enrichment(mut self, enrichment: VocabularyEnrichment) -> Self {
118 if !enrichment.is_empty() {
119 self.enrichment = Some(enrichment);
120 }
121 self
122 }
123
124 #[must_use]
126 pub fn with_embedder(mut self, embedder: Arc<dyn Embedder + Send + Sync>) -> Self {
127 self.embedder = Some(embedder);
128 self
129 }
130
131 pub fn append(&self, record: &EpisodicRecord) -> Result<()> {
133 self.append_batch(std::slice::from_ref(record))
134 }
135
136 pub fn append_batch(&self, records: &[EpisodicRecord]) -> Result<()> {
139 if records.is_empty() {
140 return Ok(());
141 }
142 let serialized = records
143 .iter()
144 .map(|record| {
145 rmp_serde::to_vec(record)
146 .map(|value| (record, value))
147 .map_err(|e| CoreError::Memory(format!("episodic serialize failed: {e}")))
148 })
149 .collect::<Result<Vec<_>>>()?;
150 let mut tx = self
151 .env
152 .begin_rw_txn()
153 .map_err(|e| CoreError::Memory(format!("episodic rw_txn failed: {e}")))?;
154 for (record, value) in &serialized {
155 match tx.put(
156 self.db,
157 record.id.as_bytes(),
158 value,
159 WriteFlags::NO_OVERWRITE,
160 ) {
161 Ok(()) => {}
162 Err(lmdb::Error::KeyExist) => {
163 tx.abort();
164 return Err(CoreError::InvalidArgs(format!(
165 "episodic record {} already exists",
166 record.id
167 )));
168 }
169 Err(e) => {
170 tx.abort();
171 return Err(CoreError::Memory(format!("episodic append failed: {e}")));
172 }
173 }
174 }
175 #[cfg(test)]
176 commit_boundary_test_hook(CommitBoundary::BeforeRawCommit, records)?;
177 tx.commit()
178 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
179 #[cfg(test)]
180 commit_boundary_test_hook(CommitBoundary::AfterRawCommit, records)?;
181 self.mutation_count
182 .fetch_add(records.len() as u64, std::sync::atomic::Ordering::Relaxed);
183 self.index_records(records)?;
187 self.clear_term_cache();
188 Ok(())
189 }
190
191 fn term_postings(&self, term: &str) -> Result<Vec<EpisodicId>> {
198 if let Ok(cache) = self.term_cache.read() {
199 if let Some(ids) = cache.get(term) {
200 return Ok(ids.clone());
201 }
202 }
203
204 let tx = self
205 .env
206 .begin_ro_txn()
207 .map_err(|e| CoreError::Memory(format!("episodic index ro_txn failed: {e}")))?;
208 let term_key = term.to_string();
209 let mut ids: Vec<EpisodicId> = Vec::new();
210 {
211 match tx.get(self.term_db, &term_key) {
214 Ok(_) => {}
215 Err(lmdb::Error::NotFound) => {
216 tx.commit().map_err(|e| {
217 CoreError::Memory(format!("episodic index commit failed: {e}"))
218 })?;
219 if let Ok(mut cache) = self.term_cache.write() {
220 cache.insert(term_key, Vec::new());
221 }
222 return Ok(ids);
223 }
224 Err(e) => {
225 return Err(CoreError::Memory(format!(
226 "episodic index read failed: {e}"
227 )));
228 }
229 }
230 let mut cursor = tx
231 .open_ro_cursor(self.term_db)
232 .map_err(|e| CoreError::Memory(format!("episodic index cursor failed: {e}")))?;
233 for (key, value) in cursor.iter_from(term_key.as_bytes()) {
234 if key != term_key.as_bytes() {
235 break;
236 }
237 if value.len() == std::mem::size_of::<EpisodicId>() {
238 if let Ok(id) = EpisodicId::from_slice(value) {
239 ids.push(id);
240 }
241 }
242 }
243 }
244 tx.commit()
245 .map_err(|e| CoreError::Memory(format!("episodic index commit failed: {e}")))?;
246 if let Ok(mut cache) = self.term_cache.write() {
247 cache.insert(term_key, ids.clone());
248 }
249 Ok(ids)
250 }
251
252 fn clear_term_cache(&self) {
253 if let Ok(mut cache) = self.term_cache.write() {
254 cache.clear();
255 }
256 }
257
258 fn index_records(&self, records: &[EpisodicRecord]) -> Result<()> {
259 let public: Vec<&EpisodicRecord> = records
260 .iter()
261 .filter(|record| !record.is_private && !record.model_exclude)
262 .collect();
263 if public.is_empty() {
264 return Ok(());
265 }
266 let mut pending: HashMap<String, Vec<&EpisodicRecord>> = HashMap::new();
267 for record in &public {
268 let base_terms = index_terms_with_aliases(&record.content, self.aliases.as_ref());
269 let enriched: Vec<String> = if let Some(ref enrichment) = self.enrichment {
270 let mut all = base_terms.clone();
271 let extra = enrichment.enrich(&base_terms);
272 all.extend(extra);
273 all.sort();
274 all.dedup();
275 all
276 } else {
277 base_terms
278 };
279 for term in enriched {
280 pending.entry(term).or_default().push(record);
281 }
282 }
283 let mut tx = self
284 .env
285 .begin_rw_txn()
286 .map_err(|e| CoreError::Memory(format!("episodic index rw_txn failed: {e}")))?;
287 for (term, records_for_term) in pending {
288 for record in records_for_term {
289 if let Err(e) = tx.put(
293 self.term_db,
294 &term.as_bytes().to_vec(),
295 &record.id.as_bytes(),
296 WriteFlags::default(),
297 ) {
298 tx.abort();
299 return Err(CoreError::Memory(format!(
300 "episodic term index write failed: {e}"
301 )));
302 }
303 }
304 }
305 tx.commit()
306 .map_err(|e| CoreError::Memory(format!("episodic term index commit failed: {e}")))?;
307 Ok(())
308 }
309
310 pub fn rebuild_sidecar(&self) -> Result<usize> {
317 let records = self.scan(None, usize::MAX)?;
318 let mut indexed = 0usize;
319 for chunk in records.chunks(5_000) {
320 self.index_records(chunk)?;
321 indexed += chunk.len();
322 }
323 self.clear_term_cache();
324 Ok(indexed)
325 }
326
327 pub fn sidecar_is_empty(&self) -> Result<bool> {
334 let tx = self
335 .env
336 .begin_ro_txn()
337 .map_err(|e| CoreError::Memory(format!("episodic index ro_txn failed: {e}")))?;
338 let mut cursor = tx
339 .open_ro_cursor(self.term_db)
340 .map_err(|e| CoreError::Memory(format!("episodic index cursor failed: {e}")))?;
341 Ok(cursor.iter().next().is_none())
342 }
343
344 pub fn record_count(&self) -> Result<u64> {
346 let tx = self
347 .env
348 .begin_ro_txn()
349 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
350 let mut cursor = tx
351 .open_ro_cursor(self.db)
352 .map_err(|e| CoreError::Memory(format!("episodic cursor failed: {e}")))?;
353 let mut count = 0u64;
354 for _ in cursor.iter() {
355 count += 1;
356 }
357 Ok(count)
358 }
359
360 pub fn append_explicit(
362 &self,
363 record: &EpisodicRecord,
364 policy: EpisodicCapturePolicy,
365 ) -> Result<bool> {
366 let prepared = record
367 .clone()
368 .with_content(policy.prepare_content(&record.content));
369 self.append(&prepared)?;
370 Ok(true)
371 }
372
373 pub fn append_explicit_batch(
375 &self,
376 records: &[EpisodicRecord],
377 policy: EpisodicCapturePolicy,
378 ) -> Result<usize> {
379 if records.is_empty() {
380 return Ok(0);
381 }
382 let prepared: Vec<EpisodicRecord> = records
383 .iter()
384 .map(|record| {
385 record
386 .clone()
387 .with_content(policy.prepare_content(&record.content))
388 })
389 .collect();
390 self.append_batch(&prepared)?;
391 Ok(prepared.len())
392 }
393
394 pub fn get(&self, id: EpisodicId) -> Result<Option<EpisodicRecord>> {
396 let tx = self
397 .env
398 .begin_ro_txn()
399 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
400 let result = tx.get(self.db, id.as_bytes());
401 match result {
402 Ok(bytes) => {
403 let record: EpisodicRecord = rmp_serde::from_slice(bytes)
404 .map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
405 tx.commit()
406 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
407 Ok(Some(record))
408 }
409 Err(lmdb::Error::NotFound) => {
410 tx.commit()
411 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
412 Ok(None)
413 }
414 Err(e) => Err(CoreError::Memory(format!("episodic get failed: {e}"))),
415 }
416 }
417
418 pub fn transition(&self, id: EpisodicId, transition: MemoryTransition) -> Result<()> {
420 let mut tx = self
421 .env
422 .begin_rw_txn()
423 .map_err(|e| CoreError::Memory(format!("episodic rw_txn failed: {e}")))?;
424 let bytes = match tx.get(self.db, id.as_bytes()) {
425 Ok(bytes) => bytes,
426 Err(lmdb::Error::NotFound) => {
427 tx.abort();
428 return Err(CoreError::InvalidArgs(format!(
429 "episodic record {id} does not exist"
430 )));
431 }
432 Err(e) => {
433 tx.abort();
434 return Err(CoreError::Memory(format!("episodic get failed: {e}")));
435 }
436 };
437 let mut record: EpisodicRecord = rmp_serde::from_slice(bytes)
438 .map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
439 record
440 .transition(transition)
441 .map_err(|e| CoreError::InvalidArgs(format!("episodic transition rejected: {e}")))?;
442 let value = rmp_serde::to_vec(&record)
443 .map_err(|e| CoreError::Memory(format!("episodic serialize failed: {e}")))?;
444 tx.put(self.db, id.as_bytes(), &value, WriteFlags::default())
445 .map_err(|e| CoreError::Memory(format!("episodic transition write failed: {e}")))?;
446 tx.commit()
447 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
448 self.mutation_count
449 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
450 Ok(())
451 }
452
453 pub fn scan(
455 &self,
456 session_id: Option<uuid::Uuid>,
457 limit: usize,
458 ) -> Result<Vec<EpisodicRecord>> {
459 if limit == 0 {
460 return Ok(Vec::new());
461 }
462 let tx = self
463 .env
464 .begin_ro_txn()
465 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
466 let mut cursor = tx
467 .open_ro_cursor(self.db)
468 .map_err(|e| CoreError::Memory(format!("episodic cursor failed: {e}")))?;
469 let mut records = Vec::new();
470 for item in cursor.iter() {
471 let (_, bytes) = item;
472 let record: EpisodicRecord = rmp_serde::from_slice(bytes)
473 .map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
474 if session_id.is_none_or(|id| record.session_id == Some(id)) {
475 records.push(record);
476 }
477 }
478 drop(cursor);
479 tx.commit()
480 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
481 records.sort_by_key(|record| (record.sequence, record.created_at, record.id));
482 records.truncate(limit);
483 Ok(records)
484 }
485
486 pub fn search(
496 &self,
497 query: &str,
498 limit: usize,
499 include_historical: bool,
500 ) -> Result<Vec<EpisodicSearchResult>> {
501 self.search_with_limits(query, limit, limit.saturating_mul(2), include_historical)
502 }
503
504 pub fn search_with_limits(
512 &self,
513 query: &str,
514 limit: usize,
515 candidate_limit: usize,
516 include_historical: bool,
517 ) -> Result<Vec<EpisodicSearchResult>> {
518 if limit == 0 {
519 return Ok(Vec::new());
520 }
521 let mut results = self.search_scored(query, limit, candidate_limit, include_historical)?;
522 if is_current_query(query) {
523 self.resolve_current(&mut results);
524 }
525 results.truncate(limit);
526 Ok(results)
527 }
528
529 fn search_scored(
534 &self,
535 query: &str,
536 limit: usize,
537 candidate_limit: usize,
538 include_historical: bool,
539 ) -> Result<Vec<EpisodicSearchResult>> {
540 let plan = QueryPlan::plan(query, limit);
541 let candidate_limit = candidate_limit.max(plan.candidate_limit);
542 let query_terms = tokenize(query);
543 let query_keys = key_index_terms_with_aliases(query, self.aliases.as_ref());
544 if query_terms.is_empty() && query_keys.is_empty() || limit == 0 {
545 return Ok(Vec::new());
546 }
547 let mut candidate_scores: HashMap<EpisodicId, usize> = HashMap::new();
548 for term in query_terms.iter().chain(query_keys.iter()) {
549 for id in self.term_postings(term)? {
550 *candidate_scores.entry(id).or_default() += 1;
551 }
552 }
553
554 let records = if candidate_scores.is_empty() {
555 if !matches!(self.sidecar_is_empty(), Ok(true)) {
560 return Ok(Vec::new());
561 }
562 self.scan(None, usize::MAX)?
564 } else {
565 let mut ranked_candidates: Vec<(EpisodicId, usize)> =
566 candidate_scores.into_iter().collect();
567 ranked_candidates.sort_by(|(left_id, left_count), (right_id, right_count)| {
568 right_count
569 .cmp(left_count)
570 .then_with(|| left_id.cmp(right_id))
571 });
572 ranked_candidates.truncate(candidate_limit);
573 self.load_records(
574 &ranked_candidates
575 .into_iter()
576 .map(|(id, _)| id)
577 .collect::<Vec<_>>(),
578 )?
579 };
580
581 let mut results = Vec::new();
582 for record in records {
583 if !include_historical && !matches!(record.validity, ValidityState::Active) {
584 continue;
585 }
586 let content_terms = tokenize(&record.content);
587 let content_keys = key_index_terms_with_aliases(&record.content, self.aliases.as_ref());
588 let reverse_map: HashMap<&String, Vec<String>> =
593 if let Some(ref enrichment) = self.enrichment {
594 if matches!(record.kind, EpisodicKind::UserStatement) {
595 query_terms
596 .iter()
597 .map(|qt| (qt, enrichment.reverse_enrich(qt)))
598 .collect()
599 } else {
600 HashMap::new()
601 }
602 } else {
603 HashMap::new()
604 };
605 let mut reverse_match_count = 0usize;
606 let matched_terms = query_terms
607 .iter()
608 .filter(|term| {
609 if content_terms.iter().any(|candidate| candidate == *term)
610 || content_keys.iter().any(|candidate| candidate == *term)
611 {
612 return true;
613 }
614 if let Some(reverse_terms) = reverse_map.get(term) {
617 let found = reverse_terms.iter().any(|rt| {
618 content_terms.iter().any(|candidate| candidate == rt)
619 || content_keys.iter().any(|candidate| candidate == rt)
620 });
621 if found {
622 reverse_match_count += 1;
623 }
624 return found;
625 }
626 false
627 })
628 .count();
629 let matched_keys = query_keys
630 .iter()
631 .filter(|term| {
632 content_keys.iter().any(|candidate| candidate == *term)
633 || content_terms.iter().any(|candidate| candidate == *term)
634 })
635 .count();
636 if matched_terms == 0 && matched_keys == 0 {
637 continue;
638 }
639 let key_bonus = if query_keys.is_empty() {
640 0.0
641 } else {
642 matched_keys as f32 / query_keys.len() as f32 * plan.key_weight
643 };
644 let role_boost = match record.kind {
645 EpisodicKind::UserStatement => 0.12,
646 _ => 0.0,
647 };
648 let effective_matched = if matches!(record.kind, EpisodicKind::UserStatement) {
649 (matched_terms + 2).min(query_terms.len())
650 } else {
651 matched_terms
652 };
653 let coverage = if query_terms.is_empty() {
654 0.0
655 } else {
656 effective_matched as f32 / query_terms.len() as f32
657 };
658 let number_bonus = if plan.number_query {
659 let has_digit = content_terms
660 .iter()
661 .any(|term| term.chars().any(|c| c.is_ascii_digit()));
662 if has_digit || contains_number_word(&record.content) {
663 0.03
664 } else {
665 0.0
666 }
667 } else {
668 0.0
669 };
670 let density = matched_terms as f32 / content_terms.len().max(1) as f32;
671 results.push(EpisodicSearchResult {
672 record,
673 score: coverage
674 + key_bonus
675 + role_boost
676 + number_bonus
677 + (reverse_match_count as f32).mul_add(0.05, density * 0.03),
678 matched_terms: matched_terms.max(matched_keys),
679 });
680 }
681 let mut session_counts: HashMap<Option<uuid::Uuid>, usize> = HashMap::new();
685 for r in &results {
686 *session_counts.entry(r.record.session_id).or_default() += 1;
687 }
688 for r in &mut results {
689 let count = session_counts
690 .get(&r.record.session_id)
691 .copied()
692 .unwrap_or(1);
693 if count > 1 {
694 r.score = 0.02f32.mul_add((count - 1).min(3) as f32, r.score);
695 }
696 }
697 let mut hash_counts: HashMap<&str, usize> = HashMap::new();
703 for r in &results {
704 *hash_counts
705 .entry(r.record.content_hash.as_str())
706 .or_default() += 1;
707 }
708 let hash_boosts: HashMap<String, f32> = results
709 .iter()
710 .map(|r| {
711 let count = hash_counts
712 .get(r.record.content_hash.as_str())
713 .copied()
714 .unwrap_or(1);
715 let boost = if count > 1 {
716 0.03 * (count - 1).min(3) as f32
717 } else {
718 0.0
719 };
720 (r.record.id.to_string(), boost)
721 })
722 .collect();
723 for r in &mut results {
724 if let Some(boost) = hash_boosts.get(&r.record.id.to_string()) {
725 r.score += boost;
726 }
727 }
728 results.sort_by(|a, b| {
729 b.score
730 .partial_cmp(&a.score)
731 .unwrap_or(std::cmp::Ordering::Equal)
732 .then_with(|| b.matched_terms.cmp(&a.matched_terms))
733 .then_with(|| a.record.content.len().cmp(&b.record.content.len()))
734 .then_with(|| a.record.sequence.cmp(&b.record.sequence))
735 .then_with(|| a.record.id.cmp(&b.record.id))
736 });
737 Ok(results)
738 }
739
740 pub fn search_with_rerank(
757 &self,
758 query: &str,
759 limit: usize,
760 candidate_limit: usize,
761 include_historical: bool,
762 alpha: f32,
763 rerank_pool: usize,
764 ) -> Result<Vec<EpisodicSearchResult>> {
765 let Some(ref embedder) = self.embedder else {
766 return self.search_with_limits(query, limit, candidate_limit, include_historical);
767 };
768 if !embedder.is_available() || limit == 0 {
769 return self.search_with_limits(query, limit, candidate_limit, include_historical);
770 }
771
772 let rerank_pool = if rerank_pool == 0 {
780 limit.max(candidate_limit).min(50)
781 } else {
782 rerank_pool.max(limit).min(50)
783 };
784 let deterministic: Vec<EpisodicSearchResult> = self
787 .search_scored(query, rerank_pool, candidate_limit, include_historical)?
788 .into_iter()
789 .take(rerank_pool)
790 .collect();
791 if deterministic.is_empty() {
792 return Ok(Vec::new());
793 }
794
795 let contents: Vec<&str> = std::iter::once(query)
797 .chain(deterministic.iter().map(|r| r.record.content.as_str()))
798 .collect();
799 let embeddings = embedder.embed_batch(&contents)?;
800 if embeddings.len() != deterministic.len() + 1 {
801 return Err(CoreError::Memory(format!(
802 "embedder returned {} vectors, expected {}",
803 embeddings.len(),
804 deterministic.len() + 1
805 )));
806 }
807 let query_vec = &embeddings[0];
808 let candidate_vecs = &embeddings[1..];
809
810 if alpha >= 2.0 {
811 let protected: Vec<EpisodicSearchResult> =
816 deterministic.into_iter().take(limit).collect();
817 let cosines: Vec<f32> = protected
818 .iter()
819 .enumerate()
820 .map(|(i, _)| cosine_sim(query_vec, &candidate_vecs[i]))
821 .collect();
822 let mut order: Vec<usize> = (0..protected.len()).collect();
823 order.sort_by(|&a, &b| {
824 cosines[b]
825 .partial_cmp(&cosines[a])
826 .unwrap_or(std::cmp::Ordering::Equal)
827 .then_with(|| a.cmp(&b))
829 });
830 let mut slots: Vec<Option<EpisodicSearchResult>> =
831 protected.into_iter().map(Some).collect();
832 let reranked: Vec<EpisodicSearchResult> =
833 order.into_iter().filter_map(|i| slots[i].take()).collect();
834 Ok(reranked)
835 } else if alpha >= 1.0 {
836 let delta = 0.05;
838 let mut reranked = deterministic;
839 let cosines: Vec<f32> = candidate_vecs
840 .iter()
841 .map(|v| cosine_sim(query_vec, v))
842 .collect();
843 let n = reranked.len();
845 for _ in 0..n {
846 let mut swapped = false;
847 for i in 0..n.saturating_sub(1) {
848 let det_gap = (reranked[i].score - reranked[i + 1].score).abs();
849 if det_gap < delta && cosines[i + 1] > cosines[i] {
850 reranked.swap(i, i + 1);
851 swapped = true;
852 }
853 }
854 if !swapped {
855 break;
856 }
857 }
858 if is_current_query(query) {
859 self.resolve_current(&mut reranked);
860 }
861 reranked.truncate(limit);
862 Ok(reranked)
863 } else {
864 let max_det = deterministic
866 .iter()
867 .map(|r| r.score)
868 .fold(0.0f32, f32::max)
869 .max(1e-9);
870
871 let mut reranked: Vec<EpisodicSearchResult> = deterministic
872 .into_iter()
873 .enumerate()
874 .map(|(i, mut r)| {
875 let cosine = cosine_sim(query_vec, &candidate_vecs[i]);
876 let det_norm = r.score / max_det;
877 r.score = alpha.mul_add(det_norm, (1.0 - alpha) * cosine);
878 r
879 })
880 .collect();
881
882 reranked.sort_by(|a, b| {
883 b.score
884 .partial_cmp(&a.score)
885 .unwrap_or(std::cmp::Ordering::Equal)
886 .then_with(|| b.matched_terms.cmp(&a.matched_terms))
887 .then_with(|| a.record.content.len().cmp(&b.record.content.len()))
888 .then_with(|| a.record.sequence.cmp(&b.record.sequence))
889 .then_with(|| a.record.id.cmp(&b.record.id))
890 });
891 if is_current_query(query) {
892 self.resolve_current(&mut reranked);
893 }
894 reranked.truncate(limit);
895 Ok(reranked)
896 }
897 }
898
899 fn resolve_current(&self, results: &mut Vec<EpisodicSearchResult>) {
926 if results.len() < 2 {
927 return;
928 }
929 let mut anchors: Vec<EpisodicSearchResult> = Vec::new();
930 let mut rest: Vec<EpisodicSearchResult> = Vec::new();
931 for result in results.drain(..) {
932 let is_anchor = matches!(result.record.kind, EpisodicKind::UserStatement)
933 && contains_change_marker(&result.record.content);
934 if is_anchor {
935 anchors.push(result);
936 } else {
937 rest.push(result);
938 }
939 }
940 if anchors.is_empty() {
941 *results = rest;
943 return;
944 }
945 anchors.sort_by(|a, b| {
946 b.record
947 .created_at
948 .cmp(&a.record.created_at)
949 .then_with(|| b.record.sequence.cmp(&a.record.sequence))
950 .then_with(|| a.record.id.cmp(&b.record.id))
951 });
952 anchors.extend(rest);
953 *results = anchors;
954 }
955
956 fn load_records(&self, ids: &[EpisodicId]) -> Result<Vec<EpisodicRecord>> {
957 let tx = self
958 .env
959 .begin_ro_txn()
960 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
961 let mut records = Vec::with_capacity(ids.len());
962 for id in ids {
963 match tx.get(self.db, id.as_bytes()) {
964 Ok(bytes) => {
965 records.push(rmp_serde::from_slice(bytes).map_err(|e| {
966 CoreError::Memory(format!("episodic deserialize failed: {e}"))
967 })?);
968 }
969 Err(lmdb::Error::NotFound) => {}
970 Err(e) => return Err(CoreError::Memory(format!("episodic get failed: {e}"))),
971 }
972 }
973 tx.commit()
974 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
975 Ok(records)
976 }
977}
978
979fn index_terms_with_aliases(text: &str, aliases: Option<&AdaptiveAliases>) -> Vec<String> {
980 tokenize(text)
981 .into_iter()
982 .chain(key_index_terms_with_aliases(text, aliases))
983 .fold(Vec::new(), |mut terms, term| {
984 if !terms.contains(&term) {
985 terms.push(term);
986 }
987 terms
988 })
989}
990
991fn tokenize(text: &str) -> Vec<String> {
992 strip_stopwords(text)
993 .split(|c: char| !c.is_alphanumeric())
994 .filter(|term| term.len() > 1)
995 .map(|term| simple_stem(&term.to_ascii_lowercase()))
996 .fold(Vec::new(), |mut terms, term| {
997 if !terms.contains(&term) {
998 terms.push(term);
999 }
1000 terms
1001 })
1002}
1003
1004const CURRENT_QUERY_WORD_CUES: &[&str] = &["current", "currently", "latest", "nowadays"];
1006
1007const CURRENT_QUERY_PHRASE_CUES: &[&str] = &["these days", "right now", "at the moment"];
1009
1010#[must_use]
1017pub fn is_current_query(query: &str) -> bool {
1018 let lowered = query.to_ascii_lowercase();
1019 let has_word = lowered
1020 .split(|c: char| !c.is_alphanumeric())
1021 .any(|token| CURRENT_QUERY_WORD_CUES.contains(&token));
1022 has_word
1023 || CURRENT_QUERY_PHRASE_CUES
1024 .iter()
1025 .any(|cue| lowered.contains(cue))
1026}
1027
1028const CHANGE_MARKERS: &[&str] = &[
1033 "switched to",
1034 "switch to",
1035 "switching to",
1036 "switched from",
1037 "changed my",
1038 "change my",
1039 "changed from",
1040 "now prefer",
1041 "now i prefer",
1042 "now i'm",
1043 "now im",
1044 "no longer",
1045 "used to",
1046 "moved to",
1047 "not anymore",
1048 "instead of",
1049 "replaced",
1050 "gave up",
1051];
1052
1053fn contains_change_marker(content: &str) -> bool {
1055 let lowered = content.to_ascii_lowercase();
1056 CHANGE_MARKERS.iter().any(|marker| lowered.contains(marker))
1057}
1058
1059const CONTRADICTION_MARKERS: &[&str] = &[
1063 "no longer",
1064 "anymore",
1065 "changed my mind",
1066 "changed my",
1067 "used to",
1068 "gave up",
1069 "just a phase",
1070 "not really",
1071 "but i",
1072];
1073
1074#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1081pub struct EpisodicConflict {
1082 pub later_record: EpisodicId,
1085 pub earlier_record: EpisodicId,
1087 pub marker: String,
1089 pub shared_terms: Vec<String>,
1091 pub later_content: String,
1093 pub earlier_content: String,
1095}
1096
1097#[must_use]
1105pub fn detect_conflicts(results: &[EpisodicSearchResult]) -> Vec<EpisodicConflict> {
1106 const MAX_CONFLICTS: usize = 10;
1107 let mut conflicts: Vec<EpisodicConflict> = Vec::new();
1108 let mut seen_pairs: Vec<(EpisodicId, EpisodicId)> = Vec::new();
1109 for (i, marked) in results.iter().enumerate() {
1110 if !matches!(marked.record.kind, EpisodicKind::UserStatement) {
1111 continue;
1112 }
1113 let lowered = marked.record.content.to_ascii_lowercase();
1114 let Some(marker) = CONTRADICTION_MARKERS
1115 .iter()
1116 .find(|m| lowered.contains(*m))
1117 .copied()
1118 else {
1119 continue;
1120 };
1121 let marked_terms = tokenize(&marked.record.content);
1122 for (j, other) in results.iter().enumerate() {
1123 if i == j || !matches!(other.record.kind, EpisodicKind::UserStatement) {
1124 continue;
1125 }
1126 if other.record.content_hash == marked.record.content_hash {
1127 continue;
1128 }
1129 let other_terms = tokenize(&other.record.content);
1130 let shared: Vec<String> = marked_terms
1131 .iter()
1132 .filter(|t| other_terms.contains(t))
1133 .cloned()
1134 .collect();
1135 if shared.len() < 2 {
1136 continue;
1137 }
1138 let (later, earlier) = if (marked.record.created_at, marked.record.sequence)
1141 > (other.record.created_at, other.record.sequence)
1142 {
1143 (&marked.record, &other.record)
1144 } else {
1145 (&other.record, &marked.record)
1146 };
1147 let pair_key = if later.id < earlier.id {
1148 (later.id, earlier.id)
1149 } else {
1150 (earlier.id, later.id)
1151 };
1152 if seen_pairs.contains(&pair_key) {
1153 continue;
1154 }
1155 seen_pairs.push(pair_key);
1156 conflicts.push(EpisodicConflict {
1157 later_record: later.id,
1158 earlier_record: earlier.id,
1159 marker: marker.to_string(),
1160 shared_terms: shared,
1161 later_content: later.content.clone(),
1162 earlier_content: earlier.content.clone(),
1163 });
1164 if conflicts.len() >= MAX_CONFLICTS {
1165 return conflicts;
1166 }
1167 }
1168 }
1169 conflicts
1170}
1171
1172fn simple_stem(word: &str) -> String {
1173 if word.len() <= 3 {
1174 return word.to_string();
1175 }
1176 for suffix in ["ies", "ied", "ing", "edly", "ed", "ly", "es", "s"] {
1177 if let Some(stem) = word.strip_suffix(suffix) {
1178 if suffix == "ies" || suffix == "ied" {
1179 return format!("{stem}y");
1180 }
1181 if stem.len() >= 2 {
1182 return stem.to_string();
1183 }
1184 }
1185 }
1186 word.to_string()
1187}
1188
1189fn cosine_sim(a: &[f32], b: &[f32]) -> f32 {
1190 let dot = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum::<f32>();
1191 let norm_a = a.iter().map(|x| x * x).sum::<f32>().sqrt();
1192 let norm_b = b.iter().map(|x| x * x).sum::<f32>().sqrt();
1193 if norm_a < 1e-9 || norm_b < 1e-9 {
1194 0.0
1195 } else {
1196 dot / (norm_a * norm_b)
1197 }
1198}
1199
1200fn contains_number_word(text: &str) -> bool {
1201 const NUMBER_WORDS: &[&str] = &[
1202 "one",
1203 "two",
1204 "three",
1205 "four",
1206 "five",
1207 "six",
1208 "seven",
1209 "eight",
1210 "nine",
1211 "ten",
1212 "eleven",
1213 "twelve",
1214 "thirteen",
1215 "fourteen",
1216 "fifteen",
1217 "sixteen",
1218 "seventeen",
1219 "eighteen",
1220 "nineteen",
1221 "twenty",
1222 "thirty",
1223 "forty",
1224 "fifty",
1225 "sixty",
1226 "seventy",
1227 "eighty",
1228 "ninety",
1229 "hundred",
1230 "thousand",
1231 "million",
1232 "billion",
1233 "dozen",
1234 "couple",
1235 "half",
1236 "quarter",
1237 "double",
1238 "triple",
1239 "twice",
1240 ];
1241 for word in text.split(|c: char| !c.is_alphanumeric()) {
1242 if word.len() >= 3 && NUMBER_WORDS.iter().any(|nw| word.eq_ignore_ascii_case(nw)) {
1243 return true;
1244 }
1245 }
1246 false
1247}
1248
1249#[cfg(test)]
1250mod tests {
1251 use super::*;
1252 use crate::MemoryStore;
1253 #[cfg(unix)]
1257 use chrono::{DateTime, Utc};
1258 use tempfile::tempdir;
1259 #[cfg(unix)]
1260 use uuid::Uuid;
1261 use wm_core::{EpisodicKind, Provenance, ProvenanceSource, ValidityState};
1262
1263 fn sample_record(sequence: u64, content: &str) -> EpisodicRecord {
1264 EpisodicRecord::new(
1265 None,
1266 sequence,
1267 EpisodicKind::Observation,
1268 content,
1269 Provenance::new(ProvenanceSource::User),
1270 )
1271 }
1272
1273 fn user_statement(sequence: u64, content: &str) -> EpisodicRecord {
1274 EpisodicRecord::new(
1275 None,
1276 sequence,
1277 EpisodicKind::UserStatement,
1278 content,
1279 Provenance::new(ProvenanceSource::User),
1280 )
1281 }
1282
1283 fn assistant_response(sequence: u64, content: &str) -> EpisodicRecord {
1284 EpisodicRecord::new(
1285 None,
1286 sequence,
1287 EpisodicKind::AssistantResponse,
1288 content,
1289 Provenance::new(ProvenanceSource::Agent),
1290 )
1291 }
1292
1293 #[cfg(unix)]
1301 fn q06_record(id: u128, sequence: u64, content: &str, created_at: &str) -> EpisodicRecord {
1302 let mut record = EpisodicRecord::new(
1303 None,
1304 sequence,
1305 EpisodicKind::Observation,
1306 content,
1307 Provenance::new(ProvenanceSource::User),
1308 )
1309 .with_id(Uuid::from_u128(id));
1310 record.created_at = DateTime::parse_from_rfc3339(created_at)
1311 .unwrap()
1312 .with_timezone(&Utc);
1313 record
1314 }
1315
1316 #[cfg(unix)]
1317 fn q06_acknowledged() -> EpisodicRecord {
1318 q06_record(601, 601, "q06 acknowledged control", "2026-01-01T00:10:01Z")
1319 }
1320
1321 #[cfg(unix)]
1322 fn run_q06_child(case: &str, store_path: &std::path::Path) {
1323 use std::io::Write;
1324 let expected_uuid = std::env::var("WM_Q06_UUID").expect("WM_Q06_UUID");
1325 let uuid = Uuid::parse_str(&expected_uuid).expect("WM_Q06_UUID must parse");
1326 let (sequence, content, created_at) = match case {
1327 "before_raw_commit" => (602, "q06 precommit candidate", "2026-01-01T00:10:02Z"),
1328 "after_raw_commit" => (603, "q06 uncertain candidate", "2026-01-01T00:10:03Z"),
1329 other => panic!("q06 child: unknown case {other}"),
1330 };
1331 let record = q06_record(uuid.as_u128(), sequence, content, created_at);
1332 let store = MemoryStore::open_default(store_path).expect("q06 child store");
1333 match store.episodic().append(&record) {
1334 Ok(()) => {
1335 println!("WM_Q06_CALLER_ACK {uuid}");
1338 let _ = std::io::stdout().flush();
1339 }
1340 Err(e) => {
1341 eprintln!("q06 child append failed: {e}");
1342 std::process::exit(2);
1343 }
1344 }
1345 }
1346
1347 #[cfg(unix)]
1348 fn run_q06_killed_case(
1349 test_filter: &str,
1350 store_path: &std::path::Path,
1351 case: &str,
1352 uuid: Uuid,
1353 ) -> Vec<String> {
1354 use std::io::{BufRead, BufReader};
1355 let exe = std::env::current_exe().expect("q06 current_exe");
1356 let mut child = std::process::Command::new(exe)
1357 .arg(test_filter)
1358 .arg("--exact")
1359 .arg("--nocapture")
1360 .env("WM_Q06_CASE", case)
1361 .env("WM_Q06_STORE", store_path)
1362 .env("WM_Q06_UUID", uuid.to_string())
1363 .stdout(std::process::Stdio::piped())
1364 .stdin(std::process::Stdio::piped())
1365 .stderr(std::process::Stdio::inherit())
1366 .spawn()
1367 .expect("q06 spawn child");
1368
1369 let stdout = child.stdout.take().expect("q06 child stdout");
1370 let (tx, rx) = std::sync::mpsc::channel::<String>();
1371 let reader = std::thread::spawn(move || {
1372 for line in BufReader::new(stdout).lines() {
1373 match line {
1374 Ok(line) => {
1375 if tx.send(line).is_err() {
1376 break;
1377 }
1378 }
1379 Err(_) => break,
1380 }
1381 }
1382 });
1383
1384 let expected = format!("WM_Q06_BOUNDARY {case} {uuid}");
1385 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
1386 let mut lines = Vec::new();
1387 let mut seen = false;
1388 while !seen {
1389 let remaining = deadline.saturating_duration_since(std::time::Instant::now());
1390 if remaining.is_zero() {
1391 break;
1392 }
1393 match rx.recv_timeout(remaining) {
1394 Ok(line) if line == expected => seen = true,
1395 Ok(line) => lines.push(line),
1396 Err(_) => break,
1397 }
1398 }
1399 if !seen {
1400 let _ = child.kill();
1401 let _ = child.wait();
1402 reader.join().ok();
1403 panic!("q06 case {case}: boundary {expected:?} not observed; lines={lines:?}");
1404 }
1405
1406 child.kill().expect("q06 kill blocked child");
1407 let status = child.wait().expect("q06 reap child");
1408 assert!(
1409 !status.success(),
1410 "q06 case {case}: killed child must not exit successfully: {status:?}"
1411 );
1412 reader.join().ok();
1413 while let Ok(line) = rx.try_recv() {
1414 lines.push(line);
1415 }
1416 assert!(
1417 !lines.iter().any(|line| line.contains("WM_Q06_CALLER_ACK")),
1418 "q06 case {case}: caller acknowledgement in a killed case invalidates the experiment: {lines:?}"
1419 );
1420 lines
1421 }
1422
1423 #[cfg(unix)]
1424 #[test]
1425 fn q06_commit_boundary_sigkill_classification() {
1426 if let Ok(case) = std::env::var("WM_Q06_CASE") {
1427 let store_path = std::env::var("WM_Q06_STORE").expect("WM_Q06_STORE");
1428 run_q06_child(&case, std::path::Path::new(&store_path));
1429 return;
1430 }
1431
1432 let dir = tempdir().unwrap();
1433 let store_path = dir.path().join("lmdb");
1434 let acknowledged = q06_acknowledged();
1435
1436 {
1438 let store = MemoryStore::open_default(&store_path).unwrap();
1439 store.episodic().append(&acknowledged).unwrap();
1440 let read = store
1441 .episodic()
1442 .get(acknowledged.id)
1443 .unwrap()
1444 .expect("acknowledged record");
1445 assert_eq!(read, acknowledged);
1446 }
1447
1448 let test_filter = "episodic::tests::q06_commit_boundary_sigkill_classification";
1449
1450 let candidate_before =
1452 q06_record(602, 602, "q06 precommit candidate", "2026-01-01T00:10:02Z");
1453 run_q06_killed_case(
1454 test_filter,
1455 &store_path,
1456 "before_raw_commit",
1457 candidate_before.id,
1458 );
1459 {
1460 let store = MemoryStore::open_default(&store_path).unwrap();
1461 assert_eq!(
1462 store.episodic().get(acknowledged.id).unwrap(),
1463 Some(acknowledged.clone()),
1464 "acknowledged record must survive a pre-commit kill unchanged"
1465 );
1466 assert_eq!(
1467 store.episodic().get(candidate_before.id).unwrap(),
1468 None,
1469 "pre-commit candidate must be absent after reopen"
1470 );
1471 }
1472
1473 let candidate_after =
1476 q06_record(603, 603, "q06 uncertain candidate", "2026-01-01T00:10:03Z");
1477 run_q06_killed_case(
1478 test_filter,
1479 &store_path,
1480 "after_raw_commit",
1481 candidate_after.id,
1482 );
1483 {
1484 let store = MemoryStore::open_default(&store_path).unwrap();
1485 assert_eq!(
1486 store.episodic().get(acknowledged.id).unwrap(),
1487 Some(acknowledged),
1488 "acknowledged record must survive a post-commit kill unchanged"
1489 );
1490 assert_eq!(
1491 store.episodic().get(candidate_before.id).unwrap(),
1492 None,
1493 "pre-commit candidate must stay absent"
1494 );
1495 assert_eq!(
1496 store.episodic().get(candidate_after.id).unwrap(),
1497 Some(candidate_after),
1498 "post-commit candidate must be present and byte-equal after reopen"
1499 );
1500 }
1501 }
1502
1503 #[test]
1504 fn current_query_detection() {
1505 assert!(is_current_query("What's my current favorite coffee?"));
1506 assert!(is_current_query("What am I currently reading these days?"));
1507 assert!(is_current_query("What's the latest book I mentioned?"));
1508 assert!(is_current_query("What's my job right now?"));
1509 assert!(is_current_query("What am I eating at the moment?"));
1510 assert!(!is_current_query("What's my favorite coffee?"));
1511 assert!(!is_current_query("Where did I volunteer in February?"));
1512 assert!(!is_current_query("What did I say about the trip?"));
1513 assert!(!is_current_query("What currency did I use in Japan?"));
1515 }
1516
1517 #[test]
1518 fn protected_rerank_preserves_membership_with_a_small_pool() {
1519 use crate::embedder::Embedder;
1524 use std::sync::Arc;
1525 use std::sync::atomic::{AtomicUsize, Ordering};
1526
1527 struct CountingEmbedder(Arc<AtomicUsize>);
1528 impl Embedder for CountingEmbedder {
1529 fn embed_batch(&self, texts: &[&str]) -> wm_core::Result<Vec<Vec<f32>>> {
1530 self.0.store(texts.len(), Ordering::Relaxed);
1531 Ok(texts
1532 .iter()
1533 .map(|t| {
1534 let mut v = vec![0.0_f32; 8];
1535 for (i, b) in t.bytes().enumerate() {
1536 v[i % 8] += f32::from(b) / 255.0;
1537 }
1538 v
1539 })
1540 .collect())
1541 }
1542 fn dimension(&self) -> usize {
1543 8
1544 }
1545 fn is_available(&self) -> bool {
1546 true
1547 }
1548 fn backend_name(&self) -> &'static str {
1549 "counting-test"
1550 }
1551 }
1552
1553 let tmp = tempdir().unwrap();
1554 let store = MemoryStore::open_default(tmp.path()).unwrap();
1555 for i in 0..20 {
1556 store
1557 .episodic()
1558 .append(&user_statement(
1559 i,
1560 &format!("Record {i} discusses topic alpha beta gamma delta epsilon"),
1561 ))
1562 .unwrap();
1563 }
1564
1565 let query = "topic alpha beta gamma delta epsilon";
1566 let expected: Vec<uuid::Uuid> = store
1569 .episodic()
1570 .search_scored(query, 5, 20, false)
1571 .unwrap()
1572 .iter()
1573 .take(5)
1574 .map(|r| r.record.id)
1575 .collect();
1576 assert_eq!(expected.len(), 5);
1577
1578 let batches = Arc::new(AtomicUsize::new(0));
1579 store.set_episodic_embedder(Arc::new(CountingEmbedder(batches.clone())));
1580 let reranked = store
1581 .episodic()
1582 .search_with_rerank(query, 5, 20, false, 2.0, 5)
1583 .unwrap();
1584 assert_eq!(
1585 batches.load(Ordering::Relaxed),
1586 6,
1587 "small rerank pool must embed pool + query only"
1588 );
1589 let mut got: Vec<uuid::Uuid> = reranked.iter().map(|r| r.record.id).collect();
1590 let mut want = expected;
1591 got.sort();
1592 want.sort();
1593 assert_eq!(
1594 got, want,
1595 "protected mode must keep the deterministic candidate set (membership), only reorder it"
1596 );
1597 }
1598
1599 #[test]
1600 fn current_query_resolution_prefers_latest_statement() {
1601 let tmp = tempdir().unwrap();
1602 let store = MemoryStore::open_default(tmp.path()).unwrap();
1603 let episodic = store.episodic();
1604 episodic
1608 .append(&user_statement(1, "My favorite coffee is dark roast."))
1609 .unwrap();
1610 episodic
1611 .append(&user_statement(
1612 2,
1613 "I really love dark roast when it comes to coffee.",
1614 ))
1615 .unwrap();
1616 episodic
1617 .append(&user_statement(3, "I've been jogging lately."))
1618 .unwrap();
1619 episodic
1620 .append(&user_statement(4, "I've switched to cold brew for coffee."))
1621 .unwrap();
1622
1623 let results = episodic
1624 .search("What's my current favorite coffee?", 5, false)
1625 .unwrap();
1626 assert!(!results.is_empty());
1627 assert!(
1628 results[0].record.content.contains("cold brew"),
1629 "current query must rank the latest statement first, got: {}",
1630 results[0].record.content
1631 );
1632 }
1633
1634 #[test]
1635 fn current_query_anchors_switched_from_template() {
1636 let tmp = tempdir().unwrap();
1642 let store = MemoryStore::open_default(tmp.path()).unwrap();
1643 let episodic = store.episodic();
1644 episodic
1645 .append(&user_statement(1, "My favorite coffee is espresso."))
1646 .unwrap();
1647 episodic
1648 .append(&user_statement(
1649 2,
1650 "I've actually switched from espresso to latte for coffee.",
1651 ))
1652 .unwrap();
1653 episodic
1654 .append(&user_statement(3, "My favorite coffee is latte."))
1655 .unwrap();
1656
1657 let results = episodic
1658 .search("What's my current favorite coffee?", 5, false)
1659 .unwrap();
1660 assert!(!results.is_empty());
1661 assert!(
1662 results[0].record.content.contains("latte"),
1663 "'switched from' must anchor the current value, got: {}",
1664 results[0].record.content
1665 );
1666 }
1667
1668 #[test]
1669 fn non_current_query_keeps_score_order() {
1670 let tmp = tempdir().unwrap();
1671 let store = MemoryStore::open_default(tmp.path()).unwrap();
1672 let episodic = store.episodic();
1673 episodic
1674 .append(&user_statement(1, "My favorite coffee is dark roast."))
1675 .unwrap();
1676 episodic
1677 .append(&user_statement(2, "I've switched to cold brew for coffee."))
1678 .unwrap();
1679
1680 let results = episodic
1683 .search("What's my favorite coffee?", 5, false)
1684 .unwrap();
1685 assert!(!results.is_empty());
1686 assert!(
1687 results[0].record.content.contains("dark roast"),
1688 "non-current query must keep score order, got: {}",
1689 results[0].record.content
1690 );
1691 }
1692
1693 #[test]
1694 fn current_resolution_anchors_on_user_statements_only() {
1695 let tmp = tempdir().unwrap();
1696 let store = MemoryStore::open_default(tmp.path()).unwrap();
1697 let episodic = store.episodic();
1698 episodic
1699 .append(&user_statement(1, "My favorite coffee is dark roast."))
1700 .unwrap();
1701 episodic
1703 .append(&assistant_response(
1704 2,
1705 "Got it, dark roast is your favorite coffee!",
1706 ))
1707 .unwrap();
1708 episodic
1709 .append(&user_statement(3, "I've switched to cold brew for coffee."))
1710 .unwrap();
1711
1712 let results = episodic
1713 .search("What's my current favorite coffee?", 5, false)
1714 .unwrap();
1715 assert!(
1716 results[0].record.content.contains("cold brew"),
1717 "user statements anchor chronology, got: {}",
1718 results[0].record.content
1719 );
1720 assert_eq!(results[0].record.kind, EpisodicKind::UserStatement);
1721 }
1722
1723 #[test]
1724 fn current_query_without_change_markers_keeps_score_order() {
1725 let tmp = tempdir().unwrap();
1726 let store = MemoryStore::open_default(tmp.path()).unwrap();
1727 let episodic = store.episodic();
1728 episodic
1730 .append(&user_statement(
1731 1,
1732 "My favorite hiking trail is Eagle Ridge.",
1733 ))
1734 .unwrap();
1735 episodic
1736 .append(&user_statement(2, "I go hiking every weekend."))
1737 .unwrap();
1738
1739 let results = episodic
1740 .search("What's my current favorite hiking trail?", 5, false)
1741 .unwrap();
1742 assert!(!results.is_empty());
1743 assert!(
1744 results[0].record.content.contains("Eagle Ridge"),
1745 "no change markers → deterministic score order, got: {}",
1746 results[0].record.content
1747 );
1748 }
1749
1750 #[test]
1751 fn detect_conflicts_flags_contradiction_with_shared_topic() {
1752 let tmp = tempdir().unwrap();
1753 let store = MemoryStore::open_default(tmp.path()).unwrap();
1754 let episodic = store.episodic();
1755 episodic
1756 .append(&user_statement(
1757 1,
1758 "I'm vegetarian now. I decided to stop eating animal products.",
1759 ))
1760 .unwrap();
1761 episodic
1762 .append(&user_statement(
1763 2,
1764 "I'm not really vegetarian anymore, I eat steak now.",
1765 ))
1766 .unwrap();
1767 episodic
1768 .append(&user_statement(3, "I went hiking yesterday."))
1769 .unwrap();
1770
1771 let results = episodic
1772 .search("vegetarian steak eating", 10, false)
1773 .unwrap();
1774 let conflicts = detect_conflicts(&results);
1775 assert_eq!(
1776 conflicts.len(),
1777 1,
1778 "the vegetarian/steak pair must be flagged, got {conflicts:?}"
1779 );
1780 let conflict = &conflicts[0];
1781 assert!(conflict.later_content.contains("steak"));
1782 assert!(conflict.earlier_content.contains("animal products"));
1783 assert!(
1784 conflict.shared_terms.iter().any(|t| t == "vegetarian"),
1785 "shared terms must include the topic: {:?}",
1786 conflict.shared_terms
1787 );
1788 }
1789
1790 #[test]
1791 fn detect_conflicts_ignores_plain_statements_and_assistant_turns() {
1792 let tmp = tempdir().unwrap();
1793 let store = MemoryStore::open_default(tmp.path()).unwrap();
1794 let episodic = store.episodic();
1795 episodic
1797 .append(&user_statement(1, "My favorite coffee is dark roast."))
1798 .unwrap();
1799 episodic
1800 .append(&user_statement(2, "I love coffee with breakfast."))
1801 .unwrap();
1802 episodic
1804 .append(&assistant_response(
1805 3,
1806 "You mentioned you no longer like tea!",
1807 ))
1808 .unwrap();
1809
1810 let results = episodic.search("coffee tea breakfast", 10, false).unwrap();
1811 assert!(detect_conflicts(&results).is_empty());
1812 }
1813
1814 #[test]
1815 fn detect_conflicts_skips_identical_content() {
1816 let tmp = tempdir().unwrap();
1817 let store = MemoryStore::open_default(tmp.path()).unwrap();
1818 let episodic = store.episodic();
1819 let record = user_statement(1, "I'm vegetarian now, but I changed my mind.");
1820 let duplicate = user_statement(2, "I'm vegetarian now, but I changed my mind.");
1821 episodic.append(&record).unwrap();
1822 episodic.append(&duplicate).unwrap();
1823
1824 let results = episodic.search("vegetarian", 10, false).unwrap();
1825 assert!(detect_conflicts(&results).is_empty());
1827 }
1828
1829 #[test]
1830 fn append_get_transition_and_reopen_roundtrip() {
1831 let tmp = tempdir().unwrap();
1832 let session = uuid::Uuid::new_v4();
1833 let record = EpisodicRecord::new(
1834 Some(session),
1835 2,
1836 EpisodicKind::Decision,
1837 "use the raw episodic lane",
1838 Provenance::new(ProvenanceSource::User).with_actor("test"),
1839 );
1840 let id = record.id;
1841 {
1842 let store = MemoryStore::open_default(tmp.path()).unwrap();
1843 let episodic = store.episodic();
1844 episodic.append(&record).unwrap();
1845 assert_eq!(episodic.get(id).unwrap().unwrap(), record);
1846 episodic
1847 .transition(
1848 id,
1849 MemoryTransition::Supersede {
1850 replacement: uuid::Uuid::new_v4(),
1851 },
1852 )
1853 .unwrap();
1854 assert!(matches!(
1855 episodic.get(id).unwrap().unwrap().validity,
1856 ValidityState::Superseded { .. }
1857 ));
1858 }
1859 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1860 let records = reopened.episodic().scan(Some(session), 10).unwrap();
1861 assert_eq!(records.len(), 1);
1862 assert_eq!(records[0].id, id);
1863 }
1864
1865 #[test]
1866 fn duplicate_append_is_rejected() {
1867 let tmp = tempdir().unwrap();
1868 let store = MemoryStore::open_default(tmp.path()).unwrap();
1869 let record = sample_record(1, "once");
1870 store.episodic().append(&record).unwrap();
1871 let error = store.episodic().append(&record).unwrap_err();
1872 assert!(error.to_string().contains("already exists"));
1873 }
1874
1875 #[test]
1876 fn raw_search_returns_canonical_records_and_skips_revoked_by_default() {
1877 let tmp = tempdir().unwrap();
1878 let store = MemoryStore::open_default(tmp.path()).unwrap();
1879 let active = sample_record(1, "Rust memory retrieval");
1880 let revoked = sample_record(2, "Rust memory retrieval old");
1881 let revoked_id = revoked.id;
1882 store.episodic().append(&active).unwrap();
1883 store.episodic().append(&revoked).unwrap();
1884 store
1885 .episodic()
1886 .transition(
1887 revoked_id,
1888 MemoryTransition::Revoke {
1889 reason: "stale".into(),
1890 },
1891 )
1892 .unwrap();
1893
1894 let current = store
1895 .episodic()
1896 .search("memory retrieval", 10, false)
1897 .unwrap();
1898 assert_eq!(current.len(), 1);
1899 assert_eq!(current[0].record.id, active.id);
1900
1901 let all = store
1902 .episodic()
1903 .search("memory retrieval", 10, true)
1904 .unwrap();
1905 assert_eq!(all.len(), 2);
1906 }
1907
1908 #[test]
1909 fn append_batch_indexes_once_and_preserves_search() {
1910 let tmp = tempdir().unwrap();
1911 let store = MemoryStore::open_default(tmp.path()).unwrap();
1912 let first = sample_record(1, "Dr. Patel scheduled a follow-up appointment");
1913 let second = sample_record(2, "unrelated grocery list");
1914 let first_id = first.id;
1915 store.episodic().append_batch(&[first, second]).unwrap();
1916 let hits = store
1917 .episodic()
1918 .search("patel appointment", 10, false)
1919 .unwrap();
1920 assert_eq!(hits.len(), 1);
1921 assert_eq!(hits[0].record.id, first_id);
1922 }
1923
1924 #[test]
1925 fn rejected_append_batch_preserves_prior_state_after_reopen() {
1926 let tmp = tempdir().unwrap();
1930 let original = sample_record(1, "acknowledged original record");
1931 let rejected_new = sample_record(2, "must not survive rejected batch");
1932 {
1933 let store = MemoryStore::open_default(tmp.path()).unwrap();
1934 store.episodic().append(&original).unwrap();
1935 let error = store
1936 .episodic()
1937 .append_batch(&[rejected_new.clone(), original.clone()])
1938 .unwrap_err();
1939 assert!(error.to_string().contains("already exists"));
1940 assert_eq!(
1941 store.episodic().get(original.id).unwrap(),
1942 Some(original.clone())
1943 );
1944 assert_eq!(store.episodic().get(rejected_new.id).unwrap(), None);
1945 }
1946
1947 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1948 assert_eq!(
1949 reopened.episodic().get(original.id).unwrap(),
1950 Some(original)
1951 );
1952 assert_eq!(reopened.episodic().get(rejected_new.id).unwrap(), None);
1953 }
1954
1955 #[test]
1956 fn lost_episodic_sidecar_rebuilds_from_raw_records_after_reopen() {
1957 let tmp = tempdir().unwrap();
1962 let record = sample_record(1, "sidecar recovery preserves searchable evidence");
1963 {
1964 let store = MemoryStore::open_default(tmp.path()).unwrap();
1965 let episodic = store.episodic();
1966 episodic.append(&record).unwrap();
1967 assert!(!episodic.sidecar_is_empty().unwrap());
1968
1969 let mut tx = store.env().begin_rw_txn().unwrap();
1970 tx.clear_db(episodic.term_db).unwrap();
1971 tx.commit().unwrap();
1972 assert!(episodic.sidecar_is_empty().unwrap());
1973 }
1974
1975 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1976 let episodic = reopened.episodic();
1977 assert_eq!(episodic.get(record.id).unwrap(), Some(record.clone()));
1978 assert!(!episodic.sidecar_is_empty().unwrap());
1979 let hits = episodic
1980 .search("sidecar searchable evidence", 10, false)
1981 .unwrap();
1982 assert_eq!(hits.len(), 1);
1983 assert_eq!(hits[0].record.id, record.id);
1984 }
1985
1986 #[test]
1987 fn append_explicit_batch_redacts_and_skips_private() {
1988 let tmp = tempdir().unwrap();
1989 let store = MemoryStore::open_default(tmp.path()).unwrap();
1990 let public = sample_record(1, "api_key=supersecret rust retrieval");
1991 let private = sample_record(2, "private rust retrieval").with_visibility(true, false);
1992 let public_id = public.id;
1993 store
1994 .episodic()
1995 .append_explicit_batch(&[public, private], EpisodicCapturePolicy::explicit_only())
1996 .unwrap();
1997 let stored = store.episodic().get(public_id).unwrap().unwrap();
1998 assert!(stored.content.contains("<REDACTED>"));
1999 let hits = store
2000 .episodic()
2001 .search("rust retrieval", 10, false)
2002 .unwrap();
2003 assert_eq!(hits.len(), 1);
2004 assert_eq!(hits[0].record.id, public_id);
2005 }
2006
2007 #[test]
2008 fn typed_keys_retrieve_vocabulary_mismatch() {
2009 let tmp = tempdir().unwrap();
2010 let store = MemoryStore::open_default(tmp.path()).unwrap();
2011 let dog = sample_record(1, "My Golden Retriever loves the park");
2012 let other = sample_record(2, "I bought a yellow dress");
2013 let dog_id = dog.id;
2014 store.episodic().append_batch(&[dog, other]).unwrap();
2015 let hits = store
2016 .episodic()
2017 .search("What breed is my dog?", 5, false)
2018 .unwrap();
2019 assert_eq!(hits[0].record.id, dog_id);
2020 }
2021
2022 #[test]
2023 fn planner_boosts_temporal_date_match() {
2024 let tmp = tempdir().unwrap();
2025 let store = MemoryStore::open_default(tmp.path()).unwrap();
2026 let dated = sample_record(1, "I volunteered on February 14th at the animal shelter");
2027 let other = sample_record(2, "I volunteered at the community garden last summer");
2028 let dated_id = dated.id;
2029 store.episodic().append_batch(&[dated, other]).unwrap();
2030 let hits = store
2031 .episodic()
2032 .search("When did I volunteer at the animal shelter?", 5, false)
2033 .unwrap();
2034 assert_eq!(hits[0].record.id, dated_id);
2035 }
2036
2037 #[test]
2038 #[ignore = "manual in-process latency profile"]
2039 fn profile_ingest_and_search_latency() {
2040 fn timed_ms(label: &str, repeats: u32, mut work: impl FnMut()) {
2041 let start = std::time::Instant::now();
2042 for _ in 0..repeats {
2043 work();
2044 }
2045 let elapsed = start.elapsed();
2046 println!(
2047 "{label}: {:.3} ms (n={repeats})",
2048 elapsed.as_secs_f64() * 1000.0 / f64::from(repeats)
2049 );
2050 }
2051
2052 let search_records: Vec<EpisodicRecord> = (0..10_000)
2053 .map(|n| {
2054 sample_record(
2055 n,
2056 if n % 5 == 0 {
2057 "Rust memory retrieval benchmark item"
2058 } else {
2059 "Unrelated episodic record"
2060 },
2061 )
2062 })
2063 .collect();
2064
2065 timed_ms("append_single_1000", 1, || {
2066 let tmp = tempdir().unwrap();
2067 let store = MemoryStore::open_default(tmp.path()).unwrap();
2068 for n in 0..1_000 {
2069 store.episodic().append(&sample_record(n, "once")).unwrap();
2070 }
2071 });
2072 timed_ms("append_batch_1000", 1, || {
2073 let tmp = tempdir().unwrap();
2074 let store = MemoryStore::open_default(tmp.path()).unwrap();
2075 let records: Vec<EpisodicRecord> =
2076 (0..1_000).map(|n| sample_record(n, "once")).collect();
2077 store.episodic().append_batch(&records).unwrap();
2078 });
2079
2080 let tmp = tempdir().unwrap();
2081 {
2082 let store = MemoryStore::open_default(tmp.path()).unwrap();
2083 store.episodic().append_batch(&search_records).unwrap();
2084 }
2085 let cold = MemoryStore::open_default(tmp.path()).unwrap();
2086 timed_ms("cold_search_10000", 1, || {
2087 let hits = cold
2088 .episodic()
2089 .search("rust memory retrieval", 10, false)
2090 .unwrap();
2091 assert!(!hits.is_empty());
2092 });
2093 timed_ms("warm_search_10000", 50, || {
2094 let hits = cold
2095 .episodic()
2096 .search("rust memory retrieval", 10, false)
2097 .unwrap();
2098 assert!(!hits.is_empty());
2099 });
2100 }
2101
2102 #[test]
2108 #[ignore = "manual in-process latency profile at realistic scale"]
2109 fn profile_search_latency_25k_realistic() {
2110 const TOTAL: u64 = 25_000;
2111 const SESSIONS: u64 = 50;
2112 let topics = [
2113 "bookshelf",
2114 "guitar",
2115 "vegetarian",
2116 "portfolio",
2117 "commute",
2118 "grandmother",
2119 "chemistry",
2120 "marathon",
2121 "internship",
2122 "yoga",
2123 "spam filter",
2124 "projector",
2125 "swimming",
2126 "cousin",
2127 "bank account",
2128 "book club",
2129 "recipe",
2130 "journal subscription",
2131 "laptop",
2132 "hiking",
2133 ];
2134 let fillers = [
2135 "We discussed the plan for the weekend and agreed on the schedule.",
2136 "The meeting notes were circulated and everyone acknowledged them.",
2137 "I explained my reasoning and the group considered the proposal.",
2138 "After the presentation we reviewed the feedback together.",
2139 "She mentioned the deadline and we adjusted the timeline accordingly.",
2140 ];
2141
2142 let records: Vec<EpisodicRecord> = (0..TOTAL)
2143 .map(|n| {
2144 let session = n * SESSIONS / TOTAL;
2145 let topic = topics[(n as usize) % topics.len()];
2146 let filler = fillers[(n as usize) % fillers.len()];
2147 let content = format!(
2148 "Session {session} note {n}: my friend Alice mentioned {topic} while {filler}"
2149 );
2150 let session_id = if n % 4 == 0 {
2151 None
2152 } else {
2153 Some(uuid::Uuid::new_v4())
2154 };
2155 EpisodicRecord::new(
2156 session_id,
2157 n,
2158 if n % 3 == 0 {
2159 EpisodicKind::UserStatement
2160 } else {
2161 EpisodicKind::AssistantResponse
2162 },
2163 content,
2164 Provenance::new(ProvenanceSource::User),
2165 )
2166 })
2167 .collect();
2168
2169 let tmp = tempdir().unwrap();
2170 {
2171 let store = MemoryStore::open_default(tmp.path()).unwrap();
2172 let t = std::time::Instant::now();
2173 store.episodic().append_batch(&records).unwrap();
2174 println!("ingest 25k: {:.1} ms", t.elapsed().as_secs_f64() * 1000.0);
2175 }
2176
2177 let cold = MemoryStore::open_default(tmp.path()).unwrap();
2179 let episodic = cold.episodic();
2180 let t = std::time::Instant::now();
2181 let hits = episodic
2182 .search("guitar grandmother recipe", 10, false)
2183 .unwrap();
2184 println!(
2185 "cold_search: {:.1} ms (hits={})",
2186 t.elapsed().as_secs_f64() * 1000.0,
2187 hits.len()
2188 );
2189
2190 let queries: Vec<String> = (0..20)
2193 .map(|i| {
2194 let a = topics[i * 7 % topics.len()];
2195 let b = topics[(i * 7 + 5) % topics.len()];
2196 format!("{a} {b} weekend plan")
2197 })
2198 .collect();
2199 let start = std::time::Instant::now();
2200 for q in &queries {
2201 let hits = episodic.search(q, 10, false).unwrap();
2202 assert!(!hits.is_empty(), "no hits for {q}");
2203 }
2204 println!(
2205 "warm_search varied p50: {:.2} ms/query",
2206 start.elapsed().as_secs_f64() * 1000.0 / queries.len() as f64
2207 );
2208
2209 let t = std::time::Instant::now();
2210 let hits = episodic
2211 .search("zzzterm zzzother zzzthird", 10, false)
2212 .unwrap();
2213 println!(
2214 "no_match_query (full-scan fallback path): {:.1} ms (hits={})",
2215 t.elapsed().as_secs_f64() * 1000.0,
2216 hits.len()
2217 );
2218 }
2219
2220 #[test]
2221 fn enrichment_bridges_vocabulary_gap_for_theater() {
2222 let tmp = tempdir().unwrap();
2223 let store = MemoryStore::open_default(tmp.path()).unwrap();
2224 store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
2225 let answer = sample_record(1, "The production I attended was The Glass Menagerie");
2227 let competing = sample_record(2, "I went to a play at the local community theater");
2228 let answer_id = answer.id;
2229 store.episodic().append_batch(&[answer, competing]).unwrap();
2230 let hits = store
2231 .episodic()
2232 .search(
2233 "What play did I attend at the local community theater?",
2234 5,
2235 false,
2236 )
2237 .unwrap();
2238 assert!(hits.iter().any(|h| h.record.id == answer_id));
2241 }
2242
2243 #[test]
2244 fn enrichment_bridges_vocabulary_gap_for_shelter() {
2245 let tmp = tempdir().unwrap();
2246 let store = MemoryStore::open_default(tmp.path()).unwrap();
2247 store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
2248 let answer = sample_record(1, "I rescued a dog from the humane society last week");
2249 let other = sample_record(2, "I bought groceries at the store");
2250 let answer_id = answer.id;
2251 store.episodic().append_batch(&[answer, other]).unwrap();
2252 let hits = store
2253 .episodic()
2254 .search("When did I volunteer at the animal shelter?", 5, false)
2255 .unwrap();
2256 assert!(hits.iter().any(|h| h.record.id == answer_id));
2257 }
2258
2259 #[test]
2260 fn session_boost_favors_sessions_with_multiple_matches() {
2261 let tmp = tempdir().unwrap();
2262 let store = MemoryStore::open_default(tmp.path()).unwrap();
2263 let session_a = uuid::Uuid::new_v4();
2264 let session_b = uuid::Uuid::new_v4();
2265 let a1 = EpisodicRecord::new(
2267 Some(session_a),
2268 1,
2269 EpisodicKind::Observation,
2270 "I love hiking in the mountains",
2271 Provenance::new(ProvenanceSource::User),
2272 );
2273 let a2 = EpisodicRecord::new(
2274 Some(session_a),
2275 2,
2276 EpisodicKind::Observation,
2277 "Hiking in the mountains is great exercise",
2278 Provenance::new(ProvenanceSource::User),
2279 );
2280 let b1 = EpisodicRecord::new(
2282 Some(session_b),
2283 1,
2284 EpisodicKind::Observation,
2285 "Hiking is fun",
2286 Provenance::new(ProvenanceSource::User),
2287 );
2288 store.episodic().append_batch(&[a1, a2, b1]).unwrap();
2289 let hits = store
2290 .episodic()
2291 .search("hiking mountains", 10, false)
2292 .unwrap();
2293 let a_ranks: Vec<usize> = hits
2296 .iter()
2297 .enumerate()
2298 .filter(|(_, h)| h.record.session_id == Some(session_a))
2299 .map(|(i, _)| i)
2300 .collect();
2301 let b_rank = hits
2302 .iter()
2303 .position(|h| h.record.session_id == Some(session_b));
2304 if let Some(br) = b_rank {
2305 assert!(a_ranks.iter().all(|&ar| ar < br));
2306 }
2307 }
2308}