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 use chrono::{DateTime, Utc};
1238 use tempfile::tempdir;
1239 use uuid::Uuid;
1240 use wm_core::{EpisodicKind, Provenance, ProvenanceSource, ValidityState};
1241
1242 fn sample_record(sequence: u64, content: &str) -> EpisodicRecord {
1243 EpisodicRecord::new(
1244 None,
1245 sequence,
1246 EpisodicKind::Observation,
1247 content,
1248 Provenance::new(ProvenanceSource::User),
1249 )
1250 }
1251
1252 fn user_statement(sequence: u64, content: &str) -> EpisodicRecord {
1253 EpisodicRecord::new(
1254 None,
1255 sequence,
1256 EpisodicKind::UserStatement,
1257 content,
1258 Provenance::new(ProvenanceSource::User),
1259 )
1260 }
1261
1262 fn assistant_response(sequence: u64, content: &str) -> EpisodicRecord {
1263 EpisodicRecord::new(
1264 None,
1265 sequence,
1266 EpisodicKind::AssistantResponse,
1267 content,
1268 Provenance::new(ProvenanceSource::Agent),
1269 )
1270 }
1271
1272 #[cfg(unix)]
1280 fn q06_record(id: u128, sequence: u64, content: &str, created_at: &str) -> EpisodicRecord {
1281 let mut record = EpisodicRecord::new(
1282 None,
1283 sequence,
1284 EpisodicKind::Observation,
1285 content,
1286 Provenance::new(ProvenanceSource::User),
1287 )
1288 .with_id(Uuid::from_u128(id));
1289 record.created_at = DateTime::parse_from_rfc3339(created_at)
1290 .unwrap()
1291 .with_timezone(&Utc);
1292 record
1293 }
1294
1295 #[cfg(unix)]
1296 fn q06_acknowledged() -> EpisodicRecord {
1297 q06_record(601, 601, "q06 acknowledged control", "2026-01-01T00:10:01Z")
1298 }
1299
1300 #[cfg(unix)]
1301 fn run_q06_child(case: &str, store_path: &std::path::Path) {
1302 use std::io::Write;
1303 let expected_uuid = std::env::var("WM_Q06_UUID").expect("WM_Q06_UUID");
1304 let uuid = Uuid::parse_str(&expected_uuid).expect("WM_Q06_UUID must parse");
1305 let (sequence, content, created_at) = match case {
1306 "before_raw_commit" => (602, "q06 precommit candidate", "2026-01-01T00:10:02Z"),
1307 "after_raw_commit" => (603, "q06 uncertain candidate", "2026-01-01T00:10:03Z"),
1308 other => panic!("q06 child: unknown case {other}"),
1309 };
1310 let record = q06_record(uuid.as_u128(), sequence, content, created_at);
1311 let store = MemoryStore::open_default(store_path).expect("q06 child store");
1312 match store.episodic().append(&record) {
1313 Ok(()) => {
1314 println!("WM_Q06_CALLER_ACK {uuid}");
1317 let _ = std::io::stdout().flush();
1318 }
1319 Err(e) => {
1320 eprintln!("q06 child append failed: {e}");
1321 std::process::exit(2);
1322 }
1323 }
1324 }
1325
1326 #[cfg(unix)]
1327 fn run_q06_killed_case(
1328 test_filter: &str,
1329 store_path: &std::path::Path,
1330 case: &str,
1331 uuid: Uuid,
1332 ) -> Vec<String> {
1333 use std::io::{BufRead, BufReader};
1334 let exe = std::env::current_exe().expect("q06 current_exe");
1335 let mut child = std::process::Command::new(exe)
1336 .arg(test_filter)
1337 .arg("--exact")
1338 .arg("--nocapture")
1339 .env("WM_Q06_CASE", case)
1340 .env("WM_Q06_STORE", store_path)
1341 .env("WM_Q06_UUID", uuid.to_string())
1342 .stdout(std::process::Stdio::piped())
1343 .stdin(std::process::Stdio::piped())
1344 .stderr(std::process::Stdio::inherit())
1345 .spawn()
1346 .expect("q06 spawn child");
1347
1348 let stdout = child.stdout.take().expect("q06 child stdout");
1349 let (tx, rx) = std::sync::mpsc::channel::<String>();
1350 let reader = std::thread::spawn(move || {
1351 for line in BufReader::new(stdout).lines() {
1352 match line {
1353 Ok(line) => {
1354 if tx.send(line).is_err() {
1355 break;
1356 }
1357 }
1358 Err(_) => break,
1359 }
1360 }
1361 });
1362
1363 let expected = format!("WM_Q06_BOUNDARY {case} {uuid}");
1364 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
1365 let mut lines = Vec::new();
1366 let mut seen = false;
1367 while !seen {
1368 let remaining = deadline.saturating_duration_since(std::time::Instant::now());
1369 if remaining.is_zero() {
1370 break;
1371 }
1372 match rx.recv_timeout(remaining) {
1373 Ok(line) if line == expected => seen = true,
1374 Ok(line) => lines.push(line),
1375 Err(_) => break,
1376 }
1377 }
1378 if !seen {
1379 let _ = child.kill();
1380 let _ = child.wait();
1381 reader.join().ok();
1382 panic!("q06 case {case}: boundary {expected:?} not observed; lines={lines:?}");
1383 }
1384
1385 child.kill().expect("q06 kill blocked child");
1386 let status = child.wait().expect("q06 reap child");
1387 assert!(
1388 !status.success(),
1389 "q06 case {case}: killed child must not exit successfully: {status:?}"
1390 );
1391 reader.join().ok();
1392 while let Ok(line) = rx.try_recv() {
1393 lines.push(line);
1394 }
1395 assert!(
1396 !lines.iter().any(|line| line.contains("WM_Q06_CALLER_ACK")),
1397 "q06 case {case}: caller acknowledgement in a killed case invalidates the experiment: {lines:?}"
1398 );
1399 lines
1400 }
1401
1402 #[cfg(unix)]
1403 #[test]
1404 fn q06_commit_boundary_sigkill_classification() {
1405 if let Ok(case) = std::env::var("WM_Q06_CASE") {
1406 let store_path = std::env::var("WM_Q06_STORE").expect("WM_Q06_STORE");
1407 run_q06_child(&case, std::path::Path::new(&store_path));
1408 return;
1409 }
1410
1411 let dir = tempdir().unwrap();
1412 let store_path = dir.path().join("lmdb");
1413 let acknowledged = q06_acknowledged();
1414
1415 {
1417 let store = MemoryStore::open_default(&store_path).unwrap();
1418 store.episodic().append(&acknowledged).unwrap();
1419 let read = store
1420 .episodic()
1421 .get(acknowledged.id)
1422 .unwrap()
1423 .expect("acknowledged record");
1424 assert_eq!(read, acknowledged);
1425 }
1426
1427 let test_filter = "episodic::tests::q06_commit_boundary_sigkill_classification";
1428
1429 let candidate_before =
1431 q06_record(602, 602, "q06 precommit candidate", "2026-01-01T00:10:02Z");
1432 run_q06_killed_case(
1433 test_filter,
1434 &store_path,
1435 "before_raw_commit",
1436 candidate_before.id,
1437 );
1438 {
1439 let store = MemoryStore::open_default(&store_path).unwrap();
1440 assert_eq!(
1441 store.episodic().get(acknowledged.id).unwrap(),
1442 Some(acknowledged.clone()),
1443 "acknowledged record must survive a pre-commit kill unchanged"
1444 );
1445 assert_eq!(
1446 store.episodic().get(candidate_before.id).unwrap(),
1447 None,
1448 "pre-commit candidate must be absent after reopen"
1449 );
1450 }
1451
1452 let candidate_after =
1455 q06_record(603, 603, "q06 uncertain candidate", "2026-01-01T00:10:03Z");
1456 run_q06_killed_case(
1457 test_filter,
1458 &store_path,
1459 "after_raw_commit",
1460 candidate_after.id,
1461 );
1462 {
1463 let store = MemoryStore::open_default(&store_path).unwrap();
1464 assert_eq!(
1465 store.episodic().get(acknowledged.id).unwrap(),
1466 Some(acknowledged),
1467 "acknowledged record must survive a post-commit kill unchanged"
1468 );
1469 assert_eq!(
1470 store.episodic().get(candidate_before.id).unwrap(),
1471 None,
1472 "pre-commit candidate must stay absent"
1473 );
1474 assert_eq!(
1475 store.episodic().get(candidate_after.id).unwrap(),
1476 Some(candidate_after),
1477 "post-commit candidate must be present and byte-equal after reopen"
1478 );
1479 }
1480 }
1481
1482 #[test]
1483 fn current_query_detection() {
1484 assert!(is_current_query("What's my current favorite coffee?"));
1485 assert!(is_current_query("What am I currently reading these days?"));
1486 assert!(is_current_query("What's the latest book I mentioned?"));
1487 assert!(is_current_query("What's my job right now?"));
1488 assert!(is_current_query("What am I eating at the moment?"));
1489 assert!(!is_current_query("What's my favorite coffee?"));
1490 assert!(!is_current_query("Where did I volunteer in February?"));
1491 assert!(!is_current_query("What did I say about the trip?"));
1492 assert!(!is_current_query("What currency did I use in Japan?"));
1494 }
1495
1496 #[test]
1497 fn current_query_resolution_prefers_latest_statement() {
1498 let tmp = tempdir().unwrap();
1499 let store = MemoryStore::open_default(tmp.path()).unwrap();
1500 let episodic = store.episodic();
1501 episodic
1505 .append(&user_statement(1, "My favorite coffee is dark roast."))
1506 .unwrap();
1507 episodic
1508 .append(&user_statement(
1509 2,
1510 "I really love dark roast when it comes to coffee.",
1511 ))
1512 .unwrap();
1513 episodic
1514 .append(&user_statement(3, "I've been jogging lately."))
1515 .unwrap();
1516 episodic
1517 .append(&user_statement(4, "I've switched to cold brew for coffee."))
1518 .unwrap();
1519
1520 let results = episodic
1521 .search("What's my current favorite coffee?", 5, false)
1522 .unwrap();
1523 assert!(!results.is_empty());
1524 assert!(
1525 results[0].record.content.contains("cold brew"),
1526 "current query must rank the latest statement first, got: {}",
1527 results[0].record.content
1528 );
1529 }
1530
1531 #[test]
1532 fn current_query_anchors_switched_from_template() {
1533 let tmp = tempdir().unwrap();
1539 let store = MemoryStore::open_default(tmp.path()).unwrap();
1540 let episodic = store.episodic();
1541 episodic
1542 .append(&user_statement(1, "My favorite coffee is espresso."))
1543 .unwrap();
1544 episodic
1545 .append(&user_statement(
1546 2,
1547 "I've actually switched from espresso to latte for coffee.",
1548 ))
1549 .unwrap();
1550 episodic
1551 .append(&user_statement(3, "My favorite coffee is latte."))
1552 .unwrap();
1553
1554 let results = episodic
1555 .search("What's my current favorite coffee?", 5, false)
1556 .unwrap();
1557 assert!(!results.is_empty());
1558 assert!(
1559 results[0].record.content.contains("latte"),
1560 "'switched from' must anchor the current value, got: {}",
1561 results[0].record.content
1562 );
1563 }
1564
1565 #[test]
1566 fn non_current_query_keeps_score_order() {
1567 let tmp = tempdir().unwrap();
1568 let store = MemoryStore::open_default(tmp.path()).unwrap();
1569 let episodic = store.episodic();
1570 episodic
1571 .append(&user_statement(1, "My favorite coffee is dark roast."))
1572 .unwrap();
1573 episodic
1574 .append(&user_statement(2, "I've switched to cold brew for coffee."))
1575 .unwrap();
1576
1577 let results = episodic
1580 .search("What's my favorite coffee?", 5, false)
1581 .unwrap();
1582 assert!(!results.is_empty());
1583 assert!(
1584 results[0].record.content.contains("dark roast"),
1585 "non-current query must keep score order, got: {}",
1586 results[0].record.content
1587 );
1588 }
1589
1590 #[test]
1591 fn current_resolution_anchors_on_user_statements_only() {
1592 let tmp = tempdir().unwrap();
1593 let store = MemoryStore::open_default(tmp.path()).unwrap();
1594 let episodic = store.episodic();
1595 episodic
1596 .append(&user_statement(1, "My favorite coffee is dark roast."))
1597 .unwrap();
1598 episodic
1600 .append(&assistant_response(
1601 2,
1602 "Got it, dark roast is your favorite coffee!",
1603 ))
1604 .unwrap();
1605 episodic
1606 .append(&user_statement(3, "I've switched to cold brew for coffee."))
1607 .unwrap();
1608
1609 let results = episodic
1610 .search("What's my current favorite coffee?", 5, false)
1611 .unwrap();
1612 assert!(
1613 results[0].record.content.contains("cold brew"),
1614 "user statements anchor chronology, got: {}",
1615 results[0].record.content
1616 );
1617 assert_eq!(results[0].record.kind, EpisodicKind::UserStatement);
1618 }
1619
1620 #[test]
1621 fn current_query_without_change_markers_keeps_score_order() {
1622 let tmp = tempdir().unwrap();
1623 let store = MemoryStore::open_default(tmp.path()).unwrap();
1624 let episodic = store.episodic();
1625 episodic
1627 .append(&user_statement(
1628 1,
1629 "My favorite hiking trail is Eagle Ridge.",
1630 ))
1631 .unwrap();
1632 episodic
1633 .append(&user_statement(2, "I go hiking every weekend."))
1634 .unwrap();
1635
1636 let results = episodic
1637 .search("What's my current favorite hiking trail?", 5, false)
1638 .unwrap();
1639 assert!(!results.is_empty());
1640 assert!(
1641 results[0].record.content.contains("Eagle Ridge"),
1642 "no change markers → deterministic score order, got: {}",
1643 results[0].record.content
1644 );
1645 }
1646
1647 #[test]
1648 fn detect_conflicts_flags_contradiction_with_shared_topic() {
1649 let tmp = tempdir().unwrap();
1650 let store = MemoryStore::open_default(tmp.path()).unwrap();
1651 let episodic = store.episodic();
1652 episodic
1653 .append(&user_statement(
1654 1,
1655 "I'm vegetarian now. I decided to stop eating animal products.",
1656 ))
1657 .unwrap();
1658 episodic
1659 .append(&user_statement(
1660 2,
1661 "I'm not really vegetarian anymore, I eat steak now.",
1662 ))
1663 .unwrap();
1664 episodic
1665 .append(&user_statement(3, "I went hiking yesterday."))
1666 .unwrap();
1667
1668 let results = episodic
1669 .search("vegetarian steak eating", 10, false)
1670 .unwrap();
1671 let conflicts = detect_conflicts(&results);
1672 assert_eq!(
1673 conflicts.len(),
1674 1,
1675 "the vegetarian/steak pair must be flagged, got {conflicts:?}"
1676 );
1677 let conflict = &conflicts[0];
1678 assert!(conflict.later_content.contains("steak"));
1679 assert!(conflict.earlier_content.contains("animal products"));
1680 assert!(
1681 conflict.shared_terms.iter().any(|t| t == "vegetarian"),
1682 "shared terms must include the topic: {:?}",
1683 conflict.shared_terms
1684 );
1685 }
1686
1687 #[test]
1688 fn detect_conflicts_ignores_plain_statements_and_assistant_turns() {
1689 let tmp = tempdir().unwrap();
1690 let store = MemoryStore::open_default(tmp.path()).unwrap();
1691 let episodic = store.episodic();
1692 episodic
1694 .append(&user_statement(1, "My favorite coffee is dark roast."))
1695 .unwrap();
1696 episodic
1697 .append(&user_statement(2, "I love coffee with breakfast."))
1698 .unwrap();
1699 episodic
1701 .append(&assistant_response(
1702 3,
1703 "You mentioned you no longer like tea!",
1704 ))
1705 .unwrap();
1706
1707 let results = episodic.search("coffee tea breakfast", 10, false).unwrap();
1708 assert!(detect_conflicts(&results).is_empty());
1709 }
1710
1711 #[test]
1712 fn detect_conflicts_skips_identical_content() {
1713 let tmp = tempdir().unwrap();
1714 let store = MemoryStore::open_default(tmp.path()).unwrap();
1715 let episodic = store.episodic();
1716 let record = user_statement(1, "I'm vegetarian now, but I changed my mind.");
1717 let duplicate = user_statement(2, "I'm vegetarian now, but I changed my mind.");
1718 episodic.append(&record).unwrap();
1719 episodic.append(&duplicate).unwrap();
1720
1721 let results = episodic.search("vegetarian", 10, false).unwrap();
1722 assert!(detect_conflicts(&results).is_empty());
1724 }
1725
1726 #[test]
1727 fn append_get_transition_and_reopen_roundtrip() {
1728 let tmp = tempdir().unwrap();
1729 let session = uuid::Uuid::new_v4();
1730 let record = EpisodicRecord::new(
1731 Some(session),
1732 2,
1733 EpisodicKind::Decision,
1734 "use the raw episodic lane",
1735 Provenance::new(ProvenanceSource::User).with_actor("test"),
1736 );
1737 let id = record.id;
1738 {
1739 let store = MemoryStore::open_default(tmp.path()).unwrap();
1740 let episodic = store.episodic();
1741 episodic.append(&record).unwrap();
1742 assert_eq!(episodic.get(id).unwrap().unwrap(), record);
1743 episodic
1744 .transition(
1745 id,
1746 MemoryTransition::Supersede {
1747 replacement: uuid::Uuid::new_v4(),
1748 },
1749 )
1750 .unwrap();
1751 assert!(matches!(
1752 episodic.get(id).unwrap().unwrap().validity,
1753 ValidityState::Superseded { .. }
1754 ));
1755 }
1756 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1757 let records = reopened.episodic().scan(Some(session), 10).unwrap();
1758 assert_eq!(records.len(), 1);
1759 assert_eq!(records[0].id, id);
1760 }
1761
1762 #[test]
1763 fn duplicate_append_is_rejected() {
1764 let tmp = tempdir().unwrap();
1765 let store = MemoryStore::open_default(tmp.path()).unwrap();
1766 let record = sample_record(1, "once");
1767 store.episodic().append(&record).unwrap();
1768 let error = store.episodic().append(&record).unwrap_err();
1769 assert!(error.to_string().contains("already exists"));
1770 }
1771
1772 #[test]
1773 fn raw_search_returns_canonical_records_and_skips_revoked_by_default() {
1774 let tmp = tempdir().unwrap();
1775 let store = MemoryStore::open_default(tmp.path()).unwrap();
1776 let active = sample_record(1, "Rust memory retrieval");
1777 let revoked = sample_record(2, "Rust memory retrieval old");
1778 let revoked_id = revoked.id;
1779 store.episodic().append(&active).unwrap();
1780 store.episodic().append(&revoked).unwrap();
1781 store
1782 .episodic()
1783 .transition(
1784 revoked_id,
1785 MemoryTransition::Revoke {
1786 reason: "stale".into(),
1787 },
1788 )
1789 .unwrap();
1790
1791 let current = store
1792 .episodic()
1793 .search("memory retrieval", 10, false)
1794 .unwrap();
1795 assert_eq!(current.len(), 1);
1796 assert_eq!(current[0].record.id, active.id);
1797
1798 let all = store
1799 .episodic()
1800 .search("memory retrieval", 10, true)
1801 .unwrap();
1802 assert_eq!(all.len(), 2);
1803 }
1804
1805 #[test]
1806 fn append_batch_indexes_once_and_preserves_search() {
1807 let tmp = tempdir().unwrap();
1808 let store = MemoryStore::open_default(tmp.path()).unwrap();
1809 let first = sample_record(1, "Dr. Patel scheduled a follow-up appointment");
1810 let second = sample_record(2, "unrelated grocery list");
1811 let first_id = first.id;
1812 store.episodic().append_batch(&[first, second]).unwrap();
1813 let hits = store
1814 .episodic()
1815 .search("patel appointment", 10, false)
1816 .unwrap();
1817 assert_eq!(hits.len(), 1);
1818 assert_eq!(hits[0].record.id, first_id);
1819 }
1820
1821 #[test]
1822 fn rejected_append_batch_preserves_prior_state_after_reopen() {
1823 let tmp = tempdir().unwrap();
1827 let original = sample_record(1, "acknowledged original record");
1828 let rejected_new = sample_record(2, "must not survive rejected batch");
1829 {
1830 let store = MemoryStore::open_default(tmp.path()).unwrap();
1831 store.episodic().append(&original).unwrap();
1832 let error = store
1833 .episodic()
1834 .append_batch(&[rejected_new.clone(), original.clone()])
1835 .unwrap_err();
1836 assert!(error.to_string().contains("already exists"));
1837 assert_eq!(
1838 store.episodic().get(original.id).unwrap(),
1839 Some(original.clone())
1840 );
1841 assert_eq!(store.episodic().get(rejected_new.id).unwrap(), None);
1842 }
1843
1844 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1845 assert_eq!(
1846 reopened.episodic().get(original.id).unwrap(),
1847 Some(original)
1848 );
1849 assert_eq!(reopened.episodic().get(rejected_new.id).unwrap(), None);
1850 }
1851
1852 #[test]
1853 fn lost_episodic_sidecar_rebuilds_from_raw_records_after_reopen() {
1854 let tmp = tempdir().unwrap();
1859 let record = sample_record(1, "sidecar recovery preserves searchable evidence");
1860 {
1861 let store = MemoryStore::open_default(tmp.path()).unwrap();
1862 let episodic = store.episodic();
1863 episodic.append(&record).unwrap();
1864 assert!(!episodic.sidecar_is_empty().unwrap());
1865
1866 let mut tx = store.env().begin_rw_txn().unwrap();
1867 tx.clear_db(episodic.term_db).unwrap();
1868 tx.commit().unwrap();
1869 assert!(episodic.sidecar_is_empty().unwrap());
1870 }
1871
1872 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1873 let episodic = reopened.episodic();
1874 assert_eq!(episodic.get(record.id).unwrap(), Some(record.clone()));
1875 assert!(!episodic.sidecar_is_empty().unwrap());
1876 let hits = episodic
1877 .search("sidecar searchable evidence", 10, false)
1878 .unwrap();
1879 assert_eq!(hits.len(), 1);
1880 assert_eq!(hits[0].record.id, record.id);
1881 }
1882
1883 #[test]
1884 fn append_explicit_batch_redacts_and_skips_private() {
1885 let tmp = tempdir().unwrap();
1886 let store = MemoryStore::open_default(tmp.path()).unwrap();
1887 let public = sample_record(1, "api_key=supersecret rust retrieval");
1888 let private = sample_record(2, "private rust retrieval").with_visibility(true, false);
1889 let public_id = public.id;
1890 store
1891 .episodic()
1892 .append_explicit_batch(&[public, private], EpisodicCapturePolicy::explicit_only())
1893 .unwrap();
1894 let stored = store.episodic().get(public_id).unwrap().unwrap();
1895 assert!(stored.content.contains("<REDACTED>"));
1896 let hits = store
1897 .episodic()
1898 .search("rust retrieval", 10, false)
1899 .unwrap();
1900 assert_eq!(hits.len(), 1);
1901 assert_eq!(hits[0].record.id, public_id);
1902 }
1903
1904 #[test]
1905 fn typed_keys_retrieve_vocabulary_mismatch() {
1906 let tmp = tempdir().unwrap();
1907 let store = MemoryStore::open_default(tmp.path()).unwrap();
1908 let dog = sample_record(1, "My Golden Retriever loves the park");
1909 let other = sample_record(2, "I bought a yellow dress");
1910 let dog_id = dog.id;
1911 store.episodic().append_batch(&[dog, other]).unwrap();
1912 let hits = store
1913 .episodic()
1914 .search("What breed is my dog?", 5, false)
1915 .unwrap();
1916 assert_eq!(hits[0].record.id, dog_id);
1917 }
1918
1919 #[test]
1920 fn planner_boosts_temporal_date_match() {
1921 let tmp = tempdir().unwrap();
1922 let store = MemoryStore::open_default(tmp.path()).unwrap();
1923 let dated = sample_record(1, "I volunteered on February 14th at the animal shelter");
1924 let other = sample_record(2, "I volunteered at the community garden last summer");
1925 let dated_id = dated.id;
1926 store.episodic().append_batch(&[dated, other]).unwrap();
1927 let hits = store
1928 .episodic()
1929 .search("When did I volunteer at the animal shelter?", 5, false)
1930 .unwrap();
1931 assert_eq!(hits[0].record.id, dated_id);
1932 }
1933
1934 #[test]
1935 #[ignore = "manual in-process latency profile"]
1936 fn profile_ingest_and_search_latency() {
1937 fn timed_ms(label: &str, repeats: u32, mut work: impl FnMut()) {
1938 let start = std::time::Instant::now();
1939 for _ in 0..repeats {
1940 work();
1941 }
1942 let elapsed = start.elapsed();
1943 println!(
1944 "{label}: {:.3} ms (n={repeats})",
1945 elapsed.as_secs_f64() * 1000.0 / f64::from(repeats)
1946 );
1947 }
1948
1949 let search_records: Vec<EpisodicRecord> = (0..10_000)
1950 .map(|n| {
1951 sample_record(
1952 n,
1953 if n % 5 == 0 {
1954 "Rust memory retrieval benchmark item"
1955 } else {
1956 "Unrelated episodic record"
1957 },
1958 )
1959 })
1960 .collect();
1961
1962 timed_ms("append_single_1000", 1, || {
1963 let tmp = tempdir().unwrap();
1964 let store = MemoryStore::open_default(tmp.path()).unwrap();
1965 for n in 0..1_000 {
1966 store.episodic().append(&sample_record(n, "once")).unwrap();
1967 }
1968 });
1969 timed_ms("append_batch_1000", 1, || {
1970 let tmp = tempdir().unwrap();
1971 let store = MemoryStore::open_default(tmp.path()).unwrap();
1972 let records: Vec<EpisodicRecord> =
1973 (0..1_000).map(|n| sample_record(n, "once")).collect();
1974 store.episodic().append_batch(&records).unwrap();
1975 });
1976
1977 let tmp = tempdir().unwrap();
1978 {
1979 let store = MemoryStore::open_default(tmp.path()).unwrap();
1980 store.episodic().append_batch(&search_records).unwrap();
1981 }
1982 let cold = MemoryStore::open_default(tmp.path()).unwrap();
1983 timed_ms("cold_search_10000", 1, || {
1984 let hits = cold
1985 .episodic()
1986 .search("rust memory retrieval", 10, false)
1987 .unwrap();
1988 assert!(!hits.is_empty());
1989 });
1990 timed_ms("warm_search_10000", 50, || {
1991 let hits = cold
1992 .episodic()
1993 .search("rust memory retrieval", 10, false)
1994 .unwrap();
1995 assert!(!hits.is_empty());
1996 });
1997 }
1998
1999 #[test]
2005 #[ignore = "manual in-process latency profile at realistic scale"]
2006 fn profile_search_latency_25k_realistic() {
2007 const TOTAL: u64 = 25_000;
2008 const SESSIONS: u64 = 50;
2009 let topics = [
2010 "bookshelf",
2011 "guitar",
2012 "vegetarian",
2013 "portfolio",
2014 "commute",
2015 "grandmother",
2016 "chemistry",
2017 "marathon",
2018 "internship",
2019 "yoga",
2020 "spam filter",
2021 "projector",
2022 "swimming",
2023 "cousin",
2024 "bank account",
2025 "book club",
2026 "recipe",
2027 "journal subscription",
2028 "laptop",
2029 "hiking",
2030 ];
2031 let fillers = [
2032 "We discussed the plan for the weekend and agreed on the schedule.",
2033 "The meeting notes were circulated and everyone acknowledged them.",
2034 "I explained my reasoning and the group considered the proposal.",
2035 "After the presentation we reviewed the feedback together.",
2036 "She mentioned the deadline and we adjusted the timeline accordingly.",
2037 ];
2038
2039 let records: Vec<EpisodicRecord> = (0..TOTAL)
2040 .map(|n| {
2041 let session = n * SESSIONS / TOTAL;
2042 let topic = topics[(n as usize) % topics.len()];
2043 let filler = fillers[(n as usize) % fillers.len()];
2044 let content = format!(
2045 "Session {session} note {n}: my friend Alice mentioned {topic} while {filler}"
2046 );
2047 let session_id = if n % 4 == 0 {
2048 None
2049 } else {
2050 Some(uuid::Uuid::new_v4())
2051 };
2052 EpisodicRecord::new(
2053 session_id,
2054 n,
2055 if n % 3 == 0 {
2056 EpisodicKind::UserStatement
2057 } else {
2058 EpisodicKind::AssistantResponse
2059 },
2060 content,
2061 Provenance::new(ProvenanceSource::User),
2062 )
2063 })
2064 .collect();
2065
2066 let tmp = tempdir().unwrap();
2067 {
2068 let store = MemoryStore::open_default(tmp.path()).unwrap();
2069 let t = std::time::Instant::now();
2070 store.episodic().append_batch(&records).unwrap();
2071 println!("ingest 25k: {:.1} ms", t.elapsed().as_secs_f64() * 1000.0);
2072 }
2073
2074 let cold = MemoryStore::open_default(tmp.path()).unwrap();
2076 let episodic = cold.episodic();
2077 let t = std::time::Instant::now();
2078 let hits = episodic
2079 .search("guitar grandmother recipe", 10, false)
2080 .unwrap();
2081 println!(
2082 "cold_search: {:.1} ms (hits={})",
2083 t.elapsed().as_secs_f64() * 1000.0,
2084 hits.len()
2085 );
2086
2087 let queries: Vec<String> = (0..20)
2090 .map(|i| {
2091 let a = topics[i * 7 % topics.len()];
2092 let b = topics[(i * 7 + 5) % topics.len()];
2093 format!("{a} {b} weekend plan")
2094 })
2095 .collect();
2096 let start = std::time::Instant::now();
2097 for q in &queries {
2098 let hits = episodic.search(q, 10, false).unwrap();
2099 assert!(!hits.is_empty(), "no hits for {q}");
2100 }
2101 println!(
2102 "warm_search varied p50: {:.2} ms/query",
2103 start.elapsed().as_secs_f64() * 1000.0 / queries.len() as f64
2104 );
2105
2106 let t = std::time::Instant::now();
2107 let hits = episodic
2108 .search("zzzterm zzzother zzzthird", 10, false)
2109 .unwrap();
2110 println!(
2111 "no_match_query (full-scan fallback path): {:.1} ms (hits={})",
2112 t.elapsed().as_secs_f64() * 1000.0,
2113 hits.len()
2114 );
2115 }
2116
2117 #[test]
2118 fn enrichment_bridges_vocabulary_gap_for_theater() {
2119 let tmp = tempdir().unwrap();
2120 let store = MemoryStore::open_default(tmp.path()).unwrap();
2121 store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
2122 let answer = sample_record(1, "The production I attended was The Glass Menagerie");
2124 let competing = sample_record(2, "I went to a play at the local community theater");
2125 let answer_id = answer.id;
2126 store.episodic().append_batch(&[answer, competing]).unwrap();
2127 let hits = store
2128 .episodic()
2129 .search(
2130 "What play did I attend at the local community theater?",
2131 5,
2132 false,
2133 )
2134 .unwrap();
2135 assert!(hits.iter().any(|h| h.record.id == answer_id));
2138 }
2139
2140 #[test]
2141 fn enrichment_bridges_vocabulary_gap_for_shelter() {
2142 let tmp = tempdir().unwrap();
2143 let store = MemoryStore::open_default(tmp.path()).unwrap();
2144 store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
2145 let answer = sample_record(1, "I rescued a dog from the humane society last week");
2146 let other = sample_record(2, "I bought groceries at the store");
2147 let answer_id = answer.id;
2148 store.episodic().append_batch(&[answer, other]).unwrap();
2149 let hits = store
2150 .episodic()
2151 .search("When did I volunteer at the animal shelter?", 5, false)
2152 .unwrap();
2153 assert!(hits.iter().any(|h| h.record.id == answer_id));
2154 }
2155
2156 #[test]
2157 fn session_boost_favors_sessions_with_multiple_matches() {
2158 let tmp = tempdir().unwrap();
2159 let store = MemoryStore::open_default(tmp.path()).unwrap();
2160 let session_a = uuid::Uuid::new_v4();
2161 let session_b = uuid::Uuid::new_v4();
2162 let a1 = EpisodicRecord::new(
2164 Some(session_a),
2165 1,
2166 EpisodicKind::Observation,
2167 "I love hiking in the mountains",
2168 Provenance::new(ProvenanceSource::User),
2169 );
2170 let a2 = EpisodicRecord::new(
2171 Some(session_a),
2172 2,
2173 EpisodicKind::Observation,
2174 "Hiking in the mountains is great exercise",
2175 Provenance::new(ProvenanceSource::User),
2176 );
2177 let b1 = EpisodicRecord::new(
2179 Some(session_b),
2180 1,
2181 EpisodicKind::Observation,
2182 "Hiking is fun",
2183 Provenance::new(ProvenanceSource::User),
2184 );
2185 store.episodic().append_batch(&[a1, a2, b1]).unwrap();
2186 let hits = store
2187 .episodic()
2188 .search("hiking mountains", 10, false)
2189 .unwrap();
2190 let a_ranks: Vec<usize> = hits
2193 .iter()
2194 .enumerate()
2195 .filter(|(_, h)| h.record.session_id == Some(session_a))
2196 .map(|(i, _)| i)
2197 .collect();
2198 let b_rank = hits
2199 .iter()
2200 .position(|h| h.record.session_id == Some(session_b));
2201 if let Some(br) = b_rank {
2202 assert!(a_ranks.iter().all(|&ar| ar < br));
2203 }
2204 }
2205}