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 ) -> Result<Vec<EpisodicSearchResult>> {
764 let Some(ref embedder) = self.embedder else {
765 return self.search_with_limits(query, limit, candidate_limit, include_historical);
766 };
767 if !embedder.is_available() || limit == 0 {
768 return self.search_with_limits(query, limit, candidate_limit, include_historical);
769 }
770
771 let rerank_pool = limit.max(candidate_limit).min(50);
773 let deterministic =
774 self.search_scored(query, rerank_pool, rerank_pool, include_historical)?;
775 if deterministic.is_empty() {
776 return Ok(Vec::new());
777 }
778
779 let contents: Vec<&str> = std::iter::once(query)
781 .chain(deterministic.iter().map(|r| r.record.content.as_str()))
782 .collect();
783 let embeddings = embedder.embed_batch(&contents)?;
784 if embeddings.len() != deterministic.len() + 1 {
785 return Err(CoreError::Memory(format!(
786 "embedder returned {} vectors, expected {}",
787 embeddings.len(),
788 deterministic.len() + 1
789 )));
790 }
791 let query_vec = &embeddings[0];
792 let candidate_vecs = &embeddings[1..];
793
794 if alpha >= 2.0 {
795 let protected: Vec<EpisodicSearchResult> =
800 deterministic.into_iter().take(limit).collect();
801 let cosines: Vec<f32> = protected
802 .iter()
803 .enumerate()
804 .map(|(i, _)| cosine_sim(query_vec, &candidate_vecs[i]))
805 .collect();
806 let mut order: Vec<usize> = (0..protected.len()).collect();
807 order.sort_by(|&a, &b| {
808 cosines[b]
809 .partial_cmp(&cosines[a])
810 .unwrap_or(std::cmp::Ordering::Equal)
811 .then_with(|| a.cmp(&b))
813 });
814 let mut slots: Vec<Option<EpisodicSearchResult>> =
815 protected.into_iter().map(Some).collect();
816 let reranked: Vec<EpisodicSearchResult> =
817 order.into_iter().filter_map(|i| slots[i].take()).collect();
818 Ok(reranked)
819 } else if alpha >= 1.0 {
820 let delta = 0.05;
822 let mut reranked = deterministic;
823 let cosines: Vec<f32> = candidate_vecs
824 .iter()
825 .map(|v| cosine_sim(query_vec, v))
826 .collect();
827 let n = reranked.len();
829 for _ in 0..n {
830 let mut swapped = false;
831 for i in 0..n.saturating_sub(1) {
832 let det_gap = (reranked[i].score - reranked[i + 1].score).abs();
833 if det_gap < delta && cosines[i + 1] > cosines[i] {
834 reranked.swap(i, i + 1);
835 swapped = true;
836 }
837 }
838 if !swapped {
839 break;
840 }
841 }
842 if is_current_query(query) {
843 self.resolve_current(&mut reranked);
844 }
845 reranked.truncate(limit);
846 Ok(reranked)
847 } else {
848 let max_det = deterministic
850 .iter()
851 .map(|r| r.score)
852 .fold(0.0f32, f32::max)
853 .max(1e-9);
854
855 let mut reranked: Vec<EpisodicSearchResult> = deterministic
856 .into_iter()
857 .enumerate()
858 .map(|(i, mut r)| {
859 let cosine = cosine_sim(query_vec, &candidate_vecs[i]);
860 let det_norm = r.score / max_det;
861 r.score = alpha.mul_add(det_norm, (1.0 - alpha) * cosine);
862 r
863 })
864 .collect();
865
866 reranked.sort_by(|a, b| {
867 b.score
868 .partial_cmp(&a.score)
869 .unwrap_or(std::cmp::Ordering::Equal)
870 .then_with(|| b.matched_terms.cmp(&a.matched_terms))
871 .then_with(|| a.record.content.len().cmp(&b.record.content.len()))
872 .then_with(|| a.record.sequence.cmp(&b.record.sequence))
873 .then_with(|| a.record.id.cmp(&b.record.id))
874 });
875 if is_current_query(query) {
876 self.resolve_current(&mut reranked);
877 }
878 reranked.truncate(limit);
879 Ok(reranked)
880 }
881 }
882
883 fn resolve_current(&self, results: &mut Vec<EpisodicSearchResult>) {
910 if results.len() < 2 {
911 return;
912 }
913 let mut anchors: Vec<EpisodicSearchResult> = Vec::new();
914 let mut rest: Vec<EpisodicSearchResult> = Vec::new();
915 for result in results.drain(..) {
916 let is_anchor = matches!(result.record.kind, EpisodicKind::UserStatement)
917 && contains_change_marker(&result.record.content);
918 if is_anchor {
919 anchors.push(result);
920 } else {
921 rest.push(result);
922 }
923 }
924 if anchors.is_empty() {
925 *results = rest;
927 return;
928 }
929 anchors.sort_by(|a, b| {
930 b.record
931 .created_at
932 .cmp(&a.record.created_at)
933 .then_with(|| b.record.sequence.cmp(&a.record.sequence))
934 .then_with(|| a.record.id.cmp(&b.record.id))
935 });
936 anchors.extend(rest);
937 *results = anchors;
938 }
939
940 fn load_records(&self, ids: &[EpisodicId]) -> Result<Vec<EpisodicRecord>> {
941 let tx = self
942 .env
943 .begin_ro_txn()
944 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
945 let mut records = Vec::with_capacity(ids.len());
946 for id in ids {
947 match tx.get(self.db, id.as_bytes()) {
948 Ok(bytes) => {
949 records.push(rmp_serde::from_slice(bytes).map_err(|e| {
950 CoreError::Memory(format!("episodic deserialize failed: {e}"))
951 })?);
952 }
953 Err(lmdb::Error::NotFound) => {}
954 Err(e) => return Err(CoreError::Memory(format!("episodic get failed: {e}"))),
955 }
956 }
957 tx.commit()
958 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
959 Ok(records)
960 }
961}
962
963fn index_terms_with_aliases(text: &str, aliases: Option<&AdaptiveAliases>) -> Vec<String> {
964 tokenize(text)
965 .into_iter()
966 .chain(key_index_terms_with_aliases(text, aliases))
967 .fold(Vec::new(), |mut terms, term| {
968 if !terms.contains(&term) {
969 terms.push(term);
970 }
971 terms
972 })
973}
974
975fn tokenize(text: &str) -> Vec<String> {
976 strip_stopwords(text)
977 .split(|c: char| !c.is_alphanumeric())
978 .filter(|term| term.len() > 1)
979 .map(|term| simple_stem(&term.to_ascii_lowercase()))
980 .fold(Vec::new(), |mut terms, term| {
981 if !terms.contains(&term) {
982 terms.push(term);
983 }
984 terms
985 })
986}
987
988const CURRENT_QUERY_WORD_CUES: &[&str] = &["current", "currently", "latest", "nowadays"];
990
991const CURRENT_QUERY_PHRASE_CUES: &[&str] = &["these days", "right now", "at the moment"];
993
994#[must_use]
1001pub fn is_current_query(query: &str) -> bool {
1002 let lowered = query.to_ascii_lowercase();
1003 let has_word = lowered
1004 .split(|c: char| !c.is_alphanumeric())
1005 .any(|token| CURRENT_QUERY_WORD_CUES.contains(&token));
1006 has_word
1007 || CURRENT_QUERY_PHRASE_CUES
1008 .iter()
1009 .any(|cue| lowered.contains(cue))
1010}
1011
1012const CHANGE_MARKERS: &[&str] = &[
1017 "switched to",
1018 "switch to",
1019 "switching to",
1020 "switched from",
1021 "changed my",
1022 "change my",
1023 "changed from",
1024 "now prefer",
1025 "now i prefer",
1026 "now i'm",
1027 "now im",
1028 "no longer",
1029 "used to",
1030 "moved to",
1031 "not anymore",
1032 "instead of",
1033 "replaced",
1034 "gave up",
1035];
1036
1037fn contains_change_marker(content: &str) -> bool {
1039 let lowered = content.to_ascii_lowercase();
1040 CHANGE_MARKERS.iter().any(|marker| lowered.contains(marker))
1041}
1042
1043const CONTRADICTION_MARKERS: &[&str] = &[
1047 "no longer",
1048 "anymore",
1049 "changed my mind",
1050 "changed my",
1051 "used to",
1052 "gave up",
1053 "just a phase",
1054 "not really",
1055 "but i",
1056];
1057
1058#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1065pub struct EpisodicConflict {
1066 pub later_record: EpisodicId,
1069 pub earlier_record: EpisodicId,
1071 pub marker: String,
1073 pub shared_terms: Vec<String>,
1075 pub later_content: String,
1077 pub earlier_content: String,
1079}
1080
1081#[must_use]
1089pub fn detect_conflicts(results: &[EpisodicSearchResult]) -> Vec<EpisodicConflict> {
1090 const MAX_CONFLICTS: usize = 10;
1091 let mut conflicts: Vec<EpisodicConflict> = Vec::new();
1092 let mut seen_pairs: Vec<(EpisodicId, EpisodicId)> = Vec::new();
1093 for (i, marked) in results.iter().enumerate() {
1094 if !matches!(marked.record.kind, EpisodicKind::UserStatement) {
1095 continue;
1096 }
1097 let lowered = marked.record.content.to_ascii_lowercase();
1098 let Some(marker) = CONTRADICTION_MARKERS
1099 .iter()
1100 .find(|m| lowered.contains(*m))
1101 .copied()
1102 else {
1103 continue;
1104 };
1105 let marked_terms = tokenize(&marked.record.content);
1106 for (j, other) in results.iter().enumerate() {
1107 if i == j || !matches!(other.record.kind, EpisodicKind::UserStatement) {
1108 continue;
1109 }
1110 if other.record.content_hash == marked.record.content_hash {
1111 continue;
1112 }
1113 let other_terms = tokenize(&other.record.content);
1114 let shared: Vec<String> = marked_terms
1115 .iter()
1116 .filter(|t| other_terms.contains(t))
1117 .cloned()
1118 .collect();
1119 if shared.len() < 2 {
1120 continue;
1121 }
1122 let (later, earlier) = if (marked.record.created_at, marked.record.sequence)
1125 > (other.record.created_at, other.record.sequence)
1126 {
1127 (&marked.record, &other.record)
1128 } else {
1129 (&other.record, &marked.record)
1130 };
1131 let pair_key = if later.id < earlier.id {
1132 (later.id, earlier.id)
1133 } else {
1134 (earlier.id, later.id)
1135 };
1136 if seen_pairs.contains(&pair_key) {
1137 continue;
1138 }
1139 seen_pairs.push(pair_key);
1140 conflicts.push(EpisodicConflict {
1141 later_record: later.id,
1142 earlier_record: earlier.id,
1143 marker: marker.to_string(),
1144 shared_terms: shared,
1145 later_content: later.content.clone(),
1146 earlier_content: earlier.content.clone(),
1147 });
1148 if conflicts.len() >= MAX_CONFLICTS {
1149 return conflicts;
1150 }
1151 }
1152 }
1153 conflicts
1154}
1155
1156fn simple_stem(word: &str) -> String {
1157 if word.len() <= 3 {
1158 return word.to_string();
1159 }
1160 for suffix in ["ies", "ied", "ing", "edly", "ed", "ly", "es", "s"] {
1161 if let Some(stem) = word.strip_suffix(suffix) {
1162 if suffix == "ies" || suffix == "ied" {
1163 return format!("{stem}y");
1164 }
1165 if stem.len() >= 2 {
1166 return stem.to_string();
1167 }
1168 }
1169 }
1170 word.to_string()
1171}
1172
1173fn cosine_sim(a: &[f32], b: &[f32]) -> f32 {
1174 let dot = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum::<f32>();
1175 let norm_a = a.iter().map(|x| x * x).sum::<f32>().sqrt();
1176 let norm_b = b.iter().map(|x| x * x).sum::<f32>().sqrt();
1177 if norm_a < 1e-9 || norm_b < 1e-9 {
1178 0.0
1179 } else {
1180 dot / (norm_a * norm_b)
1181 }
1182}
1183
1184fn contains_number_word(text: &str) -> bool {
1185 const NUMBER_WORDS: &[&str] = &[
1186 "one",
1187 "two",
1188 "three",
1189 "four",
1190 "five",
1191 "six",
1192 "seven",
1193 "eight",
1194 "nine",
1195 "ten",
1196 "eleven",
1197 "twelve",
1198 "thirteen",
1199 "fourteen",
1200 "fifteen",
1201 "sixteen",
1202 "seventeen",
1203 "eighteen",
1204 "nineteen",
1205 "twenty",
1206 "thirty",
1207 "forty",
1208 "fifty",
1209 "sixty",
1210 "seventy",
1211 "eighty",
1212 "ninety",
1213 "hundred",
1214 "thousand",
1215 "million",
1216 "billion",
1217 "dozen",
1218 "couple",
1219 "half",
1220 "quarter",
1221 "double",
1222 "triple",
1223 "twice",
1224 ];
1225 for word in text.split(|c: char| !c.is_alphanumeric()) {
1226 if word.len() >= 3 && NUMBER_WORDS.iter().any(|nw| word.eq_ignore_ascii_case(nw)) {
1227 return true;
1228 }
1229 }
1230 false
1231}
1232
1233#[cfg(test)]
1234mod tests {
1235 use super::*;
1236 use crate::MemoryStore;
1237 #[cfg(unix)]
1241 use chrono::{DateTime, Utc};
1242 use tempfile::tempdir;
1243 #[cfg(unix)]
1244 use uuid::Uuid;
1245 use wm_core::{EpisodicKind, Provenance, ProvenanceSource, ValidityState};
1246
1247 fn sample_record(sequence: u64, content: &str) -> EpisodicRecord {
1248 EpisodicRecord::new(
1249 None,
1250 sequence,
1251 EpisodicKind::Observation,
1252 content,
1253 Provenance::new(ProvenanceSource::User),
1254 )
1255 }
1256
1257 fn user_statement(sequence: u64, content: &str) -> EpisodicRecord {
1258 EpisodicRecord::new(
1259 None,
1260 sequence,
1261 EpisodicKind::UserStatement,
1262 content,
1263 Provenance::new(ProvenanceSource::User),
1264 )
1265 }
1266
1267 fn assistant_response(sequence: u64, content: &str) -> EpisodicRecord {
1268 EpisodicRecord::new(
1269 None,
1270 sequence,
1271 EpisodicKind::AssistantResponse,
1272 content,
1273 Provenance::new(ProvenanceSource::Agent),
1274 )
1275 }
1276
1277 #[cfg(unix)]
1285 fn q06_record(id: u128, sequence: u64, content: &str, created_at: &str) -> EpisodicRecord {
1286 let mut record = EpisodicRecord::new(
1287 None,
1288 sequence,
1289 EpisodicKind::Observation,
1290 content,
1291 Provenance::new(ProvenanceSource::User),
1292 )
1293 .with_id(Uuid::from_u128(id));
1294 record.created_at = DateTime::parse_from_rfc3339(created_at)
1295 .unwrap()
1296 .with_timezone(&Utc);
1297 record
1298 }
1299
1300 #[cfg(unix)]
1301 fn q06_acknowledged() -> EpisodicRecord {
1302 q06_record(601, 601, "q06 acknowledged control", "2026-01-01T00:10:01Z")
1303 }
1304
1305 #[cfg(unix)]
1306 fn run_q06_child(case: &str, store_path: &std::path::Path) {
1307 use std::io::Write;
1308 let expected_uuid = std::env::var("WM_Q06_UUID").expect("WM_Q06_UUID");
1309 let uuid = Uuid::parse_str(&expected_uuid).expect("WM_Q06_UUID must parse");
1310 let (sequence, content, created_at) = match case {
1311 "before_raw_commit" => (602, "q06 precommit candidate", "2026-01-01T00:10:02Z"),
1312 "after_raw_commit" => (603, "q06 uncertain candidate", "2026-01-01T00:10:03Z"),
1313 other => panic!("q06 child: unknown case {other}"),
1314 };
1315 let record = q06_record(uuid.as_u128(), sequence, content, created_at);
1316 let store = MemoryStore::open_default(store_path).expect("q06 child store");
1317 match store.episodic().append(&record) {
1318 Ok(()) => {
1319 println!("WM_Q06_CALLER_ACK {uuid}");
1322 let _ = std::io::stdout().flush();
1323 }
1324 Err(e) => {
1325 eprintln!("q06 child append failed: {e}");
1326 std::process::exit(2);
1327 }
1328 }
1329 }
1330
1331 #[cfg(unix)]
1332 fn run_q06_killed_case(
1333 test_filter: &str,
1334 store_path: &std::path::Path,
1335 case: &str,
1336 uuid: Uuid,
1337 ) -> Vec<String> {
1338 use std::io::{BufRead, BufReader};
1339 let exe = std::env::current_exe().expect("q06 current_exe");
1340 let mut child = std::process::Command::new(exe)
1341 .arg(test_filter)
1342 .arg("--exact")
1343 .arg("--nocapture")
1344 .env("WM_Q06_CASE", case)
1345 .env("WM_Q06_STORE", store_path)
1346 .env("WM_Q06_UUID", uuid.to_string())
1347 .stdout(std::process::Stdio::piped())
1348 .stdin(std::process::Stdio::piped())
1349 .stderr(std::process::Stdio::inherit())
1350 .spawn()
1351 .expect("q06 spawn child");
1352
1353 let stdout = child.stdout.take().expect("q06 child stdout");
1354 let (tx, rx) = std::sync::mpsc::channel::<String>();
1355 let reader = std::thread::spawn(move || {
1356 for line in BufReader::new(stdout).lines() {
1357 match line {
1358 Ok(line) => {
1359 if tx.send(line).is_err() {
1360 break;
1361 }
1362 }
1363 Err(_) => break,
1364 }
1365 }
1366 });
1367
1368 let expected = format!("WM_Q06_BOUNDARY {case} {uuid}");
1369 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
1370 let mut lines = Vec::new();
1371 let mut seen = false;
1372 while !seen {
1373 let remaining = deadline.saturating_duration_since(std::time::Instant::now());
1374 if remaining.is_zero() {
1375 break;
1376 }
1377 match rx.recv_timeout(remaining) {
1378 Ok(line) if line == expected => seen = true,
1379 Ok(line) => lines.push(line),
1380 Err(_) => break,
1381 }
1382 }
1383 if !seen {
1384 let _ = child.kill();
1385 let _ = child.wait();
1386 reader.join().ok();
1387 panic!("q06 case {case}: boundary {expected:?} not observed; lines={lines:?}");
1388 }
1389
1390 child.kill().expect("q06 kill blocked child");
1391 let status = child.wait().expect("q06 reap child");
1392 assert!(
1393 !status.success(),
1394 "q06 case {case}: killed child must not exit successfully: {status:?}"
1395 );
1396 reader.join().ok();
1397 while let Ok(line) = rx.try_recv() {
1398 lines.push(line);
1399 }
1400 assert!(
1401 !lines.iter().any(|line| line.contains("WM_Q06_CALLER_ACK")),
1402 "q06 case {case}: caller acknowledgement in a killed case invalidates the experiment: {lines:?}"
1403 );
1404 lines
1405 }
1406
1407 #[cfg(unix)]
1408 #[test]
1409 fn q06_commit_boundary_sigkill_classification() {
1410 if let Ok(case) = std::env::var("WM_Q06_CASE") {
1411 let store_path = std::env::var("WM_Q06_STORE").expect("WM_Q06_STORE");
1412 run_q06_child(&case, std::path::Path::new(&store_path));
1413 return;
1414 }
1415
1416 let dir = tempdir().unwrap();
1417 let store_path = dir.path().join("lmdb");
1418 let acknowledged = q06_acknowledged();
1419
1420 {
1422 let store = MemoryStore::open_default(&store_path).unwrap();
1423 store.episodic().append(&acknowledged).unwrap();
1424 let read = store
1425 .episodic()
1426 .get(acknowledged.id)
1427 .unwrap()
1428 .expect("acknowledged record");
1429 assert_eq!(read, acknowledged);
1430 }
1431
1432 let test_filter = "episodic::tests::q06_commit_boundary_sigkill_classification";
1433
1434 let candidate_before =
1436 q06_record(602, 602, "q06 precommit candidate", "2026-01-01T00:10:02Z");
1437 run_q06_killed_case(
1438 test_filter,
1439 &store_path,
1440 "before_raw_commit",
1441 candidate_before.id,
1442 );
1443 {
1444 let store = MemoryStore::open_default(&store_path).unwrap();
1445 assert_eq!(
1446 store.episodic().get(acknowledged.id).unwrap(),
1447 Some(acknowledged.clone()),
1448 "acknowledged record must survive a pre-commit kill unchanged"
1449 );
1450 assert_eq!(
1451 store.episodic().get(candidate_before.id).unwrap(),
1452 None,
1453 "pre-commit candidate must be absent after reopen"
1454 );
1455 }
1456
1457 let candidate_after =
1460 q06_record(603, 603, "q06 uncertain candidate", "2026-01-01T00:10:03Z");
1461 run_q06_killed_case(
1462 test_filter,
1463 &store_path,
1464 "after_raw_commit",
1465 candidate_after.id,
1466 );
1467 {
1468 let store = MemoryStore::open_default(&store_path).unwrap();
1469 assert_eq!(
1470 store.episodic().get(acknowledged.id).unwrap(),
1471 Some(acknowledged),
1472 "acknowledged record must survive a post-commit kill unchanged"
1473 );
1474 assert_eq!(
1475 store.episodic().get(candidate_before.id).unwrap(),
1476 None,
1477 "pre-commit candidate must stay absent"
1478 );
1479 assert_eq!(
1480 store.episodic().get(candidate_after.id).unwrap(),
1481 Some(candidate_after),
1482 "post-commit candidate must be present and byte-equal after reopen"
1483 );
1484 }
1485 }
1486
1487 #[test]
1488 fn current_query_detection() {
1489 assert!(is_current_query("What's my current favorite coffee?"));
1490 assert!(is_current_query("What am I currently reading these days?"));
1491 assert!(is_current_query("What's the latest book I mentioned?"));
1492 assert!(is_current_query("What's my job right now?"));
1493 assert!(is_current_query("What am I eating at the moment?"));
1494 assert!(!is_current_query("What's my favorite coffee?"));
1495 assert!(!is_current_query("Where did I volunteer in February?"));
1496 assert!(!is_current_query("What did I say about the trip?"));
1497 assert!(!is_current_query("What currency did I use in Japan?"));
1499 }
1500
1501 #[test]
1502 fn current_query_resolution_prefers_latest_statement() {
1503 let tmp = tempdir().unwrap();
1504 let store = MemoryStore::open_default(tmp.path()).unwrap();
1505 let episodic = store.episodic();
1506 episodic
1510 .append(&user_statement(1, "My favorite coffee is dark roast."))
1511 .unwrap();
1512 episodic
1513 .append(&user_statement(
1514 2,
1515 "I really love dark roast when it comes to coffee.",
1516 ))
1517 .unwrap();
1518 episodic
1519 .append(&user_statement(3, "I've been jogging lately."))
1520 .unwrap();
1521 episodic
1522 .append(&user_statement(4, "I've switched to cold brew for coffee."))
1523 .unwrap();
1524
1525 let results = episodic
1526 .search("What's my current favorite coffee?", 5, false)
1527 .unwrap();
1528 assert!(!results.is_empty());
1529 assert!(
1530 results[0].record.content.contains("cold brew"),
1531 "current query must rank the latest statement first, got: {}",
1532 results[0].record.content
1533 );
1534 }
1535
1536 #[test]
1537 fn current_query_anchors_switched_from_template() {
1538 let tmp = tempdir().unwrap();
1544 let store = MemoryStore::open_default(tmp.path()).unwrap();
1545 let episodic = store.episodic();
1546 episodic
1547 .append(&user_statement(1, "My favorite coffee is espresso."))
1548 .unwrap();
1549 episodic
1550 .append(&user_statement(
1551 2,
1552 "I've actually switched from espresso to latte for coffee.",
1553 ))
1554 .unwrap();
1555 episodic
1556 .append(&user_statement(3, "My favorite coffee is latte."))
1557 .unwrap();
1558
1559 let results = episodic
1560 .search("What's my current favorite coffee?", 5, false)
1561 .unwrap();
1562 assert!(!results.is_empty());
1563 assert!(
1564 results[0].record.content.contains("latte"),
1565 "'switched from' must anchor the current value, got: {}",
1566 results[0].record.content
1567 );
1568 }
1569
1570 #[test]
1571 fn non_current_query_keeps_score_order() {
1572 let tmp = tempdir().unwrap();
1573 let store = MemoryStore::open_default(tmp.path()).unwrap();
1574 let episodic = store.episodic();
1575 episodic
1576 .append(&user_statement(1, "My favorite coffee is dark roast."))
1577 .unwrap();
1578 episodic
1579 .append(&user_statement(2, "I've switched to cold brew for coffee."))
1580 .unwrap();
1581
1582 let results = episodic
1585 .search("What's my favorite coffee?", 5, false)
1586 .unwrap();
1587 assert!(!results.is_empty());
1588 assert!(
1589 results[0].record.content.contains("dark roast"),
1590 "non-current query must keep score order, got: {}",
1591 results[0].record.content
1592 );
1593 }
1594
1595 #[test]
1596 fn current_resolution_anchors_on_user_statements_only() {
1597 let tmp = tempdir().unwrap();
1598 let store = MemoryStore::open_default(tmp.path()).unwrap();
1599 let episodic = store.episodic();
1600 episodic
1601 .append(&user_statement(1, "My favorite coffee is dark roast."))
1602 .unwrap();
1603 episodic
1605 .append(&assistant_response(
1606 2,
1607 "Got it, dark roast is your favorite coffee!",
1608 ))
1609 .unwrap();
1610 episodic
1611 .append(&user_statement(3, "I've switched to cold brew for coffee."))
1612 .unwrap();
1613
1614 let results = episodic
1615 .search("What's my current favorite coffee?", 5, false)
1616 .unwrap();
1617 assert!(
1618 results[0].record.content.contains("cold brew"),
1619 "user statements anchor chronology, got: {}",
1620 results[0].record.content
1621 );
1622 assert_eq!(results[0].record.kind, EpisodicKind::UserStatement);
1623 }
1624
1625 #[test]
1626 fn current_query_without_change_markers_keeps_score_order() {
1627 let tmp = tempdir().unwrap();
1628 let store = MemoryStore::open_default(tmp.path()).unwrap();
1629 let episodic = store.episodic();
1630 episodic
1632 .append(&user_statement(
1633 1,
1634 "My favorite hiking trail is Eagle Ridge.",
1635 ))
1636 .unwrap();
1637 episodic
1638 .append(&user_statement(2, "I go hiking every weekend."))
1639 .unwrap();
1640
1641 let results = episodic
1642 .search("What's my current favorite hiking trail?", 5, false)
1643 .unwrap();
1644 assert!(!results.is_empty());
1645 assert!(
1646 results[0].record.content.contains("Eagle Ridge"),
1647 "no change markers → deterministic score order, got: {}",
1648 results[0].record.content
1649 );
1650 }
1651
1652 #[test]
1653 fn detect_conflicts_flags_contradiction_with_shared_topic() {
1654 let tmp = tempdir().unwrap();
1655 let store = MemoryStore::open_default(tmp.path()).unwrap();
1656 let episodic = store.episodic();
1657 episodic
1658 .append(&user_statement(
1659 1,
1660 "I'm vegetarian now. I decided to stop eating animal products.",
1661 ))
1662 .unwrap();
1663 episodic
1664 .append(&user_statement(
1665 2,
1666 "I'm not really vegetarian anymore, I eat steak now.",
1667 ))
1668 .unwrap();
1669 episodic
1670 .append(&user_statement(3, "I went hiking yesterday."))
1671 .unwrap();
1672
1673 let results = episodic
1674 .search("vegetarian steak eating", 10, false)
1675 .unwrap();
1676 let conflicts = detect_conflicts(&results);
1677 assert_eq!(
1678 conflicts.len(),
1679 1,
1680 "the vegetarian/steak pair must be flagged, got {conflicts:?}"
1681 );
1682 let conflict = &conflicts[0];
1683 assert!(conflict.later_content.contains("steak"));
1684 assert!(conflict.earlier_content.contains("animal products"));
1685 assert!(
1686 conflict.shared_terms.iter().any(|t| t == "vegetarian"),
1687 "shared terms must include the topic: {:?}",
1688 conflict.shared_terms
1689 );
1690 }
1691
1692 #[test]
1693 fn detect_conflicts_ignores_plain_statements_and_assistant_turns() {
1694 let tmp = tempdir().unwrap();
1695 let store = MemoryStore::open_default(tmp.path()).unwrap();
1696 let episodic = store.episodic();
1697 episodic
1699 .append(&user_statement(1, "My favorite coffee is dark roast."))
1700 .unwrap();
1701 episodic
1702 .append(&user_statement(2, "I love coffee with breakfast."))
1703 .unwrap();
1704 episodic
1706 .append(&assistant_response(
1707 3,
1708 "You mentioned you no longer like tea!",
1709 ))
1710 .unwrap();
1711
1712 let results = episodic.search("coffee tea breakfast", 10, false).unwrap();
1713 assert!(detect_conflicts(&results).is_empty());
1714 }
1715
1716 #[test]
1717 fn detect_conflicts_skips_identical_content() {
1718 let tmp = tempdir().unwrap();
1719 let store = MemoryStore::open_default(tmp.path()).unwrap();
1720 let episodic = store.episodic();
1721 let record = user_statement(1, "I'm vegetarian now, but I changed my mind.");
1722 let duplicate = user_statement(2, "I'm vegetarian now, but I changed my mind.");
1723 episodic.append(&record).unwrap();
1724 episodic.append(&duplicate).unwrap();
1725
1726 let results = episodic.search("vegetarian", 10, false).unwrap();
1727 assert!(detect_conflicts(&results).is_empty());
1729 }
1730
1731 #[test]
1732 fn append_get_transition_and_reopen_roundtrip() {
1733 let tmp = tempdir().unwrap();
1734 let session = uuid::Uuid::new_v4();
1735 let record = EpisodicRecord::new(
1736 Some(session),
1737 2,
1738 EpisodicKind::Decision,
1739 "use the raw episodic lane",
1740 Provenance::new(ProvenanceSource::User).with_actor("test"),
1741 );
1742 let id = record.id;
1743 {
1744 let store = MemoryStore::open_default(tmp.path()).unwrap();
1745 let episodic = store.episodic();
1746 episodic.append(&record).unwrap();
1747 assert_eq!(episodic.get(id).unwrap().unwrap(), record);
1748 episodic
1749 .transition(
1750 id,
1751 MemoryTransition::Supersede {
1752 replacement: uuid::Uuid::new_v4(),
1753 },
1754 )
1755 .unwrap();
1756 assert!(matches!(
1757 episodic.get(id).unwrap().unwrap().validity,
1758 ValidityState::Superseded { .. }
1759 ));
1760 }
1761 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1762 let records = reopened.episodic().scan(Some(session), 10).unwrap();
1763 assert_eq!(records.len(), 1);
1764 assert_eq!(records[0].id, id);
1765 }
1766
1767 #[test]
1768 fn duplicate_append_is_rejected() {
1769 let tmp = tempdir().unwrap();
1770 let store = MemoryStore::open_default(tmp.path()).unwrap();
1771 let record = sample_record(1, "once");
1772 store.episodic().append(&record).unwrap();
1773 let error = store.episodic().append(&record).unwrap_err();
1774 assert!(error.to_string().contains("already exists"));
1775 }
1776
1777 #[test]
1778 fn raw_search_returns_canonical_records_and_skips_revoked_by_default() {
1779 let tmp = tempdir().unwrap();
1780 let store = MemoryStore::open_default(tmp.path()).unwrap();
1781 let active = sample_record(1, "Rust memory retrieval");
1782 let revoked = sample_record(2, "Rust memory retrieval old");
1783 let revoked_id = revoked.id;
1784 store.episodic().append(&active).unwrap();
1785 store.episodic().append(&revoked).unwrap();
1786 store
1787 .episodic()
1788 .transition(
1789 revoked_id,
1790 MemoryTransition::Revoke {
1791 reason: "stale".into(),
1792 },
1793 )
1794 .unwrap();
1795
1796 let current = store
1797 .episodic()
1798 .search("memory retrieval", 10, false)
1799 .unwrap();
1800 assert_eq!(current.len(), 1);
1801 assert_eq!(current[0].record.id, active.id);
1802
1803 let all = store
1804 .episodic()
1805 .search("memory retrieval", 10, true)
1806 .unwrap();
1807 assert_eq!(all.len(), 2);
1808 }
1809
1810 #[test]
1811 fn append_batch_indexes_once_and_preserves_search() {
1812 let tmp = tempdir().unwrap();
1813 let store = MemoryStore::open_default(tmp.path()).unwrap();
1814 let first = sample_record(1, "Dr. Patel scheduled a follow-up appointment");
1815 let second = sample_record(2, "unrelated grocery list");
1816 let first_id = first.id;
1817 store.episodic().append_batch(&[first, second]).unwrap();
1818 let hits = store
1819 .episodic()
1820 .search("patel appointment", 10, false)
1821 .unwrap();
1822 assert_eq!(hits.len(), 1);
1823 assert_eq!(hits[0].record.id, first_id);
1824 }
1825
1826 #[test]
1827 fn rejected_append_batch_preserves_prior_state_after_reopen() {
1828 let tmp = tempdir().unwrap();
1832 let original = sample_record(1, "acknowledged original record");
1833 let rejected_new = sample_record(2, "must not survive rejected batch");
1834 {
1835 let store = MemoryStore::open_default(tmp.path()).unwrap();
1836 store.episodic().append(&original).unwrap();
1837 let error = store
1838 .episodic()
1839 .append_batch(&[rejected_new.clone(), original.clone()])
1840 .unwrap_err();
1841 assert!(error.to_string().contains("already exists"));
1842 assert_eq!(
1843 store.episodic().get(original.id).unwrap(),
1844 Some(original.clone())
1845 );
1846 assert_eq!(store.episodic().get(rejected_new.id).unwrap(), None);
1847 }
1848
1849 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1850 assert_eq!(
1851 reopened.episodic().get(original.id).unwrap(),
1852 Some(original)
1853 );
1854 assert_eq!(reopened.episodic().get(rejected_new.id).unwrap(), None);
1855 }
1856
1857 #[test]
1858 fn lost_episodic_sidecar_rebuilds_from_raw_records_after_reopen() {
1859 let tmp = tempdir().unwrap();
1864 let record = sample_record(1, "sidecar recovery preserves searchable evidence");
1865 {
1866 let store = MemoryStore::open_default(tmp.path()).unwrap();
1867 let episodic = store.episodic();
1868 episodic.append(&record).unwrap();
1869 assert!(!episodic.sidecar_is_empty().unwrap());
1870
1871 let mut tx = store.env().begin_rw_txn().unwrap();
1872 tx.clear_db(episodic.term_db).unwrap();
1873 tx.commit().unwrap();
1874 assert!(episodic.sidecar_is_empty().unwrap());
1875 }
1876
1877 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1878 let episodic = reopened.episodic();
1879 assert_eq!(episodic.get(record.id).unwrap(), Some(record.clone()));
1880 assert!(!episodic.sidecar_is_empty().unwrap());
1881 let hits = episodic
1882 .search("sidecar searchable evidence", 10, false)
1883 .unwrap();
1884 assert_eq!(hits.len(), 1);
1885 assert_eq!(hits[0].record.id, record.id);
1886 }
1887
1888 #[test]
1889 fn append_explicit_batch_redacts_and_skips_private() {
1890 let tmp = tempdir().unwrap();
1891 let store = MemoryStore::open_default(tmp.path()).unwrap();
1892 let public = sample_record(1, "api_key=supersecret rust retrieval");
1893 let private = sample_record(2, "private rust retrieval").with_visibility(true, false);
1894 let public_id = public.id;
1895 store
1896 .episodic()
1897 .append_explicit_batch(&[public, private], EpisodicCapturePolicy::explicit_only())
1898 .unwrap();
1899 let stored = store.episodic().get(public_id).unwrap().unwrap();
1900 assert!(stored.content.contains("<REDACTED>"));
1901 let hits = store
1902 .episodic()
1903 .search("rust retrieval", 10, false)
1904 .unwrap();
1905 assert_eq!(hits.len(), 1);
1906 assert_eq!(hits[0].record.id, public_id);
1907 }
1908
1909 #[test]
1910 fn typed_keys_retrieve_vocabulary_mismatch() {
1911 let tmp = tempdir().unwrap();
1912 let store = MemoryStore::open_default(tmp.path()).unwrap();
1913 let dog = sample_record(1, "My Golden Retriever loves the park");
1914 let other = sample_record(2, "I bought a yellow dress");
1915 let dog_id = dog.id;
1916 store.episodic().append_batch(&[dog, other]).unwrap();
1917 let hits = store
1918 .episodic()
1919 .search("What breed is my dog?", 5, false)
1920 .unwrap();
1921 assert_eq!(hits[0].record.id, dog_id);
1922 }
1923
1924 #[test]
1925 fn planner_boosts_temporal_date_match() {
1926 let tmp = tempdir().unwrap();
1927 let store = MemoryStore::open_default(tmp.path()).unwrap();
1928 let dated = sample_record(1, "I volunteered on February 14th at the animal shelter");
1929 let other = sample_record(2, "I volunteered at the community garden last summer");
1930 let dated_id = dated.id;
1931 store.episodic().append_batch(&[dated, other]).unwrap();
1932 let hits = store
1933 .episodic()
1934 .search("When did I volunteer at the animal shelter?", 5, false)
1935 .unwrap();
1936 assert_eq!(hits[0].record.id, dated_id);
1937 }
1938
1939 #[test]
1940 #[ignore = "manual in-process latency profile"]
1941 fn profile_ingest_and_search_latency() {
1942 fn timed_ms(label: &str, repeats: u32, mut work: impl FnMut()) {
1943 let start = std::time::Instant::now();
1944 for _ in 0..repeats {
1945 work();
1946 }
1947 let elapsed = start.elapsed();
1948 println!(
1949 "{label}: {:.3} ms (n={repeats})",
1950 elapsed.as_secs_f64() * 1000.0 / f64::from(repeats)
1951 );
1952 }
1953
1954 let search_records: Vec<EpisodicRecord> = (0..10_000)
1955 .map(|n| {
1956 sample_record(
1957 n,
1958 if n % 5 == 0 {
1959 "Rust memory retrieval benchmark item"
1960 } else {
1961 "Unrelated episodic record"
1962 },
1963 )
1964 })
1965 .collect();
1966
1967 timed_ms("append_single_1000", 1, || {
1968 let tmp = tempdir().unwrap();
1969 let store = MemoryStore::open_default(tmp.path()).unwrap();
1970 for n in 0..1_000 {
1971 store.episodic().append(&sample_record(n, "once")).unwrap();
1972 }
1973 });
1974 timed_ms("append_batch_1000", 1, || {
1975 let tmp = tempdir().unwrap();
1976 let store = MemoryStore::open_default(tmp.path()).unwrap();
1977 let records: Vec<EpisodicRecord> =
1978 (0..1_000).map(|n| sample_record(n, "once")).collect();
1979 store.episodic().append_batch(&records).unwrap();
1980 });
1981
1982 let tmp = tempdir().unwrap();
1983 {
1984 let store = MemoryStore::open_default(tmp.path()).unwrap();
1985 store.episodic().append_batch(&search_records).unwrap();
1986 }
1987 let cold = MemoryStore::open_default(tmp.path()).unwrap();
1988 timed_ms("cold_search_10000", 1, || {
1989 let hits = cold
1990 .episodic()
1991 .search("rust memory retrieval", 10, false)
1992 .unwrap();
1993 assert!(!hits.is_empty());
1994 });
1995 timed_ms("warm_search_10000", 50, || {
1996 let hits = cold
1997 .episodic()
1998 .search("rust memory retrieval", 10, false)
1999 .unwrap();
2000 assert!(!hits.is_empty());
2001 });
2002 }
2003
2004 #[test]
2010 #[ignore = "manual in-process latency profile at realistic scale"]
2011 fn profile_search_latency_25k_realistic() {
2012 const TOTAL: u64 = 25_000;
2013 const SESSIONS: u64 = 50;
2014 let topics = [
2015 "bookshelf",
2016 "guitar",
2017 "vegetarian",
2018 "portfolio",
2019 "commute",
2020 "grandmother",
2021 "chemistry",
2022 "marathon",
2023 "internship",
2024 "yoga",
2025 "spam filter",
2026 "projector",
2027 "swimming",
2028 "cousin",
2029 "bank account",
2030 "book club",
2031 "recipe",
2032 "journal subscription",
2033 "laptop",
2034 "hiking",
2035 ];
2036 let fillers = [
2037 "We discussed the plan for the weekend and agreed on the schedule.",
2038 "The meeting notes were circulated and everyone acknowledged them.",
2039 "I explained my reasoning and the group considered the proposal.",
2040 "After the presentation we reviewed the feedback together.",
2041 "She mentioned the deadline and we adjusted the timeline accordingly.",
2042 ];
2043
2044 let records: Vec<EpisodicRecord> = (0..TOTAL)
2045 .map(|n| {
2046 let session = n * SESSIONS / TOTAL;
2047 let topic = topics[(n as usize) % topics.len()];
2048 let filler = fillers[(n as usize) % fillers.len()];
2049 let content = format!(
2050 "Session {session} note {n}: my friend Alice mentioned {topic} while {filler}"
2051 );
2052 let session_id = if n % 4 == 0 {
2053 None
2054 } else {
2055 Some(uuid::Uuid::new_v4())
2056 };
2057 EpisodicRecord::new(
2058 session_id,
2059 n,
2060 if n % 3 == 0 {
2061 EpisodicKind::UserStatement
2062 } else {
2063 EpisodicKind::AssistantResponse
2064 },
2065 content,
2066 Provenance::new(ProvenanceSource::User),
2067 )
2068 })
2069 .collect();
2070
2071 let tmp = tempdir().unwrap();
2072 {
2073 let store = MemoryStore::open_default(tmp.path()).unwrap();
2074 let t = std::time::Instant::now();
2075 store.episodic().append_batch(&records).unwrap();
2076 println!("ingest 25k: {:.1} ms", t.elapsed().as_secs_f64() * 1000.0);
2077 }
2078
2079 let cold = MemoryStore::open_default(tmp.path()).unwrap();
2081 let episodic = cold.episodic();
2082 let t = std::time::Instant::now();
2083 let hits = episodic
2084 .search("guitar grandmother recipe", 10, false)
2085 .unwrap();
2086 println!(
2087 "cold_search: {:.1} ms (hits={})",
2088 t.elapsed().as_secs_f64() * 1000.0,
2089 hits.len()
2090 );
2091
2092 let queries: Vec<String> = (0..20)
2095 .map(|i| {
2096 let a = topics[i * 7 % topics.len()];
2097 let b = topics[(i * 7 + 5) % topics.len()];
2098 format!("{a} {b} weekend plan")
2099 })
2100 .collect();
2101 let start = std::time::Instant::now();
2102 for q in &queries {
2103 let hits = episodic.search(q, 10, false).unwrap();
2104 assert!(!hits.is_empty(), "no hits for {q}");
2105 }
2106 println!(
2107 "warm_search varied p50: {:.2} ms/query",
2108 start.elapsed().as_secs_f64() * 1000.0 / queries.len() as f64
2109 );
2110
2111 let t = std::time::Instant::now();
2112 let hits = episodic
2113 .search("zzzterm zzzother zzzthird", 10, false)
2114 .unwrap();
2115 println!(
2116 "no_match_query (full-scan fallback path): {:.1} ms (hits={})",
2117 t.elapsed().as_secs_f64() * 1000.0,
2118 hits.len()
2119 );
2120 }
2121
2122 #[test]
2123 fn enrichment_bridges_vocabulary_gap_for_theater() {
2124 let tmp = tempdir().unwrap();
2125 let store = MemoryStore::open_default(tmp.path()).unwrap();
2126 store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
2127 let answer = sample_record(1, "The production I attended was The Glass Menagerie");
2129 let competing = sample_record(2, "I went to a play at the local community theater");
2130 let answer_id = answer.id;
2131 store.episodic().append_batch(&[answer, competing]).unwrap();
2132 let hits = store
2133 .episodic()
2134 .search(
2135 "What play did I attend at the local community theater?",
2136 5,
2137 false,
2138 )
2139 .unwrap();
2140 assert!(hits.iter().any(|h| h.record.id == answer_id));
2143 }
2144
2145 #[test]
2146 fn enrichment_bridges_vocabulary_gap_for_shelter() {
2147 let tmp = tempdir().unwrap();
2148 let store = MemoryStore::open_default(tmp.path()).unwrap();
2149 store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
2150 let answer = sample_record(1, "I rescued a dog from the humane society last week");
2151 let other = sample_record(2, "I bought groceries at the store");
2152 let answer_id = answer.id;
2153 store.episodic().append_batch(&[answer, other]).unwrap();
2154 let hits = store
2155 .episodic()
2156 .search("When did I volunteer at the animal shelter?", 5, false)
2157 .unwrap();
2158 assert!(hits.iter().any(|h| h.record.id == answer_id));
2159 }
2160
2161 #[test]
2162 fn session_boost_favors_sessions_with_multiple_matches() {
2163 let tmp = tempdir().unwrap();
2164 let store = MemoryStore::open_default(tmp.path()).unwrap();
2165 let session_a = uuid::Uuid::new_v4();
2166 let session_b = uuid::Uuid::new_v4();
2167 let a1 = EpisodicRecord::new(
2169 Some(session_a),
2170 1,
2171 EpisodicKind::Observation,
2172 "I love hiking in the mountains",
2173 Provenance::new(ProvenanceSource::User),
2174 );
2175 let a2 = EpisodicRecord::new(
2176 Some(session_a),
2177 2,
2178 EpisodicKind::Observation,
2179 "Hiking in the mountains is great exercise",
2180 Provenance::new(ProvenanceSource::User),
2181 );
2182 let b1 = EpisodicRecord::new(
2184 Some(session_b),
2185 1,
2186 EpisodicKind::Observation,
2187 "Hiking is fun",
2188 Provenance::new(ProvenanceSource::User),
2189 );
2190 store.episodic().append_batch(&[a1, a2, b1]).unwrap();
2191 let hits = store
2192 .episodic()
2193 .search("hiking mountains", 10, false)
2194 .unwrap();
2195 let a_ranks: Vec<usize> = hits
2198 .iter()
2199 .enumerate()
2200 .filter(|(_, h)| h.record.session_id == Some(session_a))
2201 .map(|(i, _)| i)
2202 .collect();
2203 let b_rank = hits
2204 .iter()
2205 .position(|h| h.record.session_id == Some(session_b));
2206 if let Some(br) = b_rank {
2207 assert!(a_ranks.iter().all(|&ar| ar < br));
2208 }
2209 }
2210}