1use lmdb::{Cursor, Database, Environment, Transaction, WriteFlags};
4use serde::{Deserialize, Serialize};
5use sha2::{Digest, Sha256};
6use std::collections::HashMap;
7use std::sync::{Arc, RwLock};
8use wm_core::{
9 CoreError, EpisodicCapturePolicy, EpisodicId, EpisodicKind, EpisodicRecord, MemoryTransition,
10 Result, ValidityState,
11};
12
13use crate::embedder::Embedder;
14use crate::enrichment::VocabularyEnrichment;
15use crate::episodic_keys::{AdaptiveAliases, key_index_terms_with_aliases};
16use crate::query_planner::QueryPlan;
17use crate::search::strip_stopwords;
18
19#[cfg(test)]
26#[derive(Clone, Copy, PartialEq, Eq)]
27enum CommitBoundary {
28 BeforeRawCommit,
29 AfterRawCommit,
30}
31
32#[cfg(test)]
33fn commit_boundary_test_hook(boundary: CommitBoundary, records: &[EpisodicRecord]) -> Result<()> {
34 use std::io::{Read, Write};
35 let Ok(case) = std::env::var("WM_Q06_CASE") else {
36 return Ok(());
37 };
38 let expected = match boundary {
39 CommitBoundary::BeforeRawCommit => "before_raw_commit",
40 CommitBoundary::AfterRawCommit => "after_raw_commit",
41 };
42 if case != expected {
43 return Ok(());
44 }
45 let uuid = std::env::var("WM_Q06_UUID")
46 .map_err(|_| CoreError::Memory("q06 hook: WM_Q06_UUID is required".into()))?;
47 let names: Vec<String> = records.iter().map(|record| record.id.to_string()).collect();
48 if names.len() != 1 || names[0] != uuid {
49 return Err(CoreError::Memory(
50 "q06 hook: candidate set must be the single WM_Q06_UUID record".into(),
51 ));
52 }
53 println!("WM_Q06_BOUNDARY {expected} {uuid}");
54 let _ = std::io::stdout().flush();
55 let mut release = [0_u8; 1];
56 match std::io::stdin().read(&mut release) {
57 Ok(0) => Err(CoreError::Memory(
58 "q06 hook: stdin closed before release".into(),
59 )),
60 Ok(_) => Ok(()),
61 Err(e) => Err(CoreError::Memory(format!(
62 "q06 hook: stdin read failed: {e}"
63 ))),
64 }
65}
66
67#[derive(Debug, Clone)]
69pub struct EpisodicSearchResult {
70 pub record: EpisodicRecord,
71 pub score: f32,
72 pub matched_terms: usize,
73}
74
75pub struct EpisodicStore<'a> {
77 env: &'a Environment,
78 db: Database,
79 term_db: Database,
80 term_cache: Arc<RwLock<HashMap<String, Vec<EpisodicId>>>>,
81 mutation_count: &'a std::sync::atomic::AtomicU64,
82 embedder: Option<Arc<dyn Embedder + Send + Sync>>,
83 aliases: Option<AdaptiveAliases>,
84 enrichment: Option<VocabularyEnrichment>,
85}
86
87impl<'a> EpisodicStore<'a> {
88 pub(crate) fn new(
89 env: &'a Environment,
90 db: Database,
91 term_db: Database,
92 term_cache: Arc<RwLock<HashMap<String, Vec<EpisodicId>>>>,
93 mutation_count: &'a std::sync::atomic::AtomicU64,
94 ) -> Self {
95 Self {
96 env,
97 db,
98 term_db,
99 term_cache,
100 mutation_count,
101 embedder: None,
102 aliases: None,
103 enrichment: None,
104 }
105 }
106
107 #[must_use]
109 pub fn with_adaptive_aliases(mut self, aliases: AdaptiveAliases) -> Self {
110 if !aliases.is_empty() {
111 self.aliases = Some(aliases);
112 }
113 self
114 }
115
116 #[must_use]
118 pub fn with_enrichment(mut self, enrichment: VocabularyEnrichment) -> Self {
119 if !enrichment.is_empty() {
120 self.enrichment = Some(enrichment);
121 }
122 self
123 }
124
125 #[must_use]
127 pub fn with_embedder(mut self, embedder: Arc<dyn Embedder + Send + Sync>) -> Self {
128 self.embedder = Some(embedder);
129 self
130 }
131
132 pub fn append(&self, record: &EpisodicRecord) -> Result<()> {
134 self.append_batch(std::slice::from_ref(record))
135 }
136
137 pub fn append_batch(&self, records: &[EpisodicRecord]) -> Result<()> {
140 if records.is_empty() {
141 return Ok(());
142 }
143 let serialized = records
144 .iter()
145 .map(|record| {
146 rmp_serde::to_vec(record)
147 .map(|value| (record, value))
148 .map_err(|e| CoreError::Memory(format!("episodic serialize failed: {e}")))
149 })
150 .collect::<Result<Vec<_>>>()?;
151 let mut tx = self
152 .env
153 .begin_rw_txn()
154 .map_err(|e| CoreError::Memory(format!("episodic rw_txn failed: {e}")))?;
155 for (record, value) in &serialized {
156 match tx.put(
157 self.db,
158 record.id.as_bytes(),
159 value,
160 WriteFlags::NO_OVERWRITE,
161 ) {
162 Ok(()) => {}
163 Err(lmdb::Error::KeyExist) => {
164 tx.abort();
165 return Err(CoreError::InvalidArgs(format!(
166 "episodic record {} already exists",
167 record.id
168 )));
169 }
170 Err(e) => {
171 tx.abort();
172 return Err(CoreError::Memory(format!("episodic append failed: {e}")));
173 }
174 }
175 }
176 #[cfg(test)]
177 commit_boundary_test_hook(CommitBoundary::BeforeRawCommit, records)?;
178 tx.commit()
179 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
180 #[cfg(test)]
181 commit_boundary_test_hook(CommitBoundary::AfterRawCommit, records)?;
182 self.mutation_count
183 .fetch_add(records.len() as u64, std::sync::atomic::Ordering::Relaxed);
184 self.index_records(records)?;
188 self.clear_term_cache();
189 Ok(())
190 }
191
192 fn term_postings(&self, term: &str) -> Result<Vec<EpisodicId>> {
199 let term = index_safe_term(term);
202 if let Ok(cache) = self.term_cache.read() {
203 if let Some(ids) = cache.get(&term) {
204 return Ok(ids.clone());
205 }
206 }
207
208 let tx = self
209 .env
210 .begin_ro_txn()
211 .map_err(|e| CoreError::Memory(format!("episodic index ro_txn failed: {e}")))?;
212 let term_key = term;
213 let mut ids: Vec<EpisodicId> = Vec::new();
214 {
215 match tx.get(self.term_db, &term_key) {
218 Ok(_) => {}
219 Err(lmdb::Error::NotFound) => {
220 tx.commit().map_err(|e| {
221 CoreError::Memory(format!("episodic index commit failed: {e}"))
222 })?;
223 if let Ok(mut cache) = self.term_cache.write() {
224 cache.insert(term_key, Vec::new());
225 }
226 return Ok(ids);
227 }
228 Err(e) => {
229 return Err(CoreError::Memory(format!(
230 "episodic index read failed: {e}"
231 )));
232 }
233 }
234 let mut cursor = tx
235 .open_ro_cursor(self.term_db)
236 .map_err(|e| CoreError::Memory(format!("episodic index cursor failed: {e}")))?;
237 for (key, value) in cursor.iter_from(term_key.as_bytes()) {
238 if key != term_key.as_bytes() {
239 break;
240 }
241 if value.len() == std::mem::size_of::<EpisodicId>() {
242 if let Ok(id) = EpisodicId::from_slice(value) {
243 ids.push(id);
244 }
245 }
246 }
247 }
248 tx.commit()
249 .map_err(|e| CoreError::Memory(format!("episodic index commit failed: {e}")))?;
250 if let Ok(mut cache) = self.term_cache.write() {
251 cache.insert(term_key, ids.clone());
252 }
253 Ok(ids)
254 }
255
256 fn clear_term_cache(&self) {
257 if let Ok(mut cache) = self.term_cache.write() {
258 cache.clear();
259 }
260 }
261
262 fn record_terms(&self, record: &EpisodicRecord) -> std::collections::BTreeSet<String> {
267 if record.is_private || record.model_exclude {
268 return std::collections::BTreeSet::new();
269 }
270 let base = index_terms_with_aliases(&record.content, self.aliases.as_ref());
271 let terms: std::collections::BTreeSet<String> =
272 if let Some(ref enrichment) = self.enrichment {
273 let mut all = base.clone();
274 all.extend(enrichment.enrich(&base));
275 all.into_iter().collect()
276 } else {
277 base.into_iter().collect()
278 };
279 terms.into_iter().map(|t| index_safe_term(&t)).collect()
282 }
283
284 fn index_records(&self, records: &[EpisodicRecord]) -> Result<()> {
285 let mut pending: HashMap<String, Vec<&EpisodicRecord>> = HashMap::new();
286 for record in records {
287 for term in self.record_terms(record) {
288 pending.entry(term).or_default().push(record);
289 }
290 }
291 if pending.is_empty() {
292 return Ok(());
293 }
294 let mut tx = self
295 .env
296 .begin_rw_txn()
297 .map_err(|e| CoreError::Memory(format!("episodic index rw_txn failed: {e}")))?;
298 for (term, records_for_term) in pending {
299 for record in records_for_term {
300 if let Err(e) = tx.put(
304 self.term_db,
305 &term,
306 &record.id.as_bytes(),
307 WriteFlags::default(),
308 ) {
309 tx.abort();
310 return Err(CoreError::Memory(format!(
311 "episodic term index write failed: {e}"
312 )));
313 }
314 }
315 }
316 tx.commit()
317 .map_err(|e| CoreError::Memory(format!("episodic term index commit failed: {e}")))?;
318 Ok(())
319 }
320
321 pub fn rebuild_sidecar(&self) -> Result<usize> {
328 let records = self.scan(None, usize::MAX)?;
329 let mut indexed = 0usize;
330 for chunk in records.chunks(5_000) {
331 self.index_records(chunk)?;
332 indexed += chunk.len();
333 }
334 self.clear_term_cache();
335 Ok(indexed)
336 }
337
338 pub fn sidecar_is_empty(&self) -> Result<bool> {
345 let tx = self
346 .env
347 .begin_ro_txn()
348 .map_err(|e| CoreError::Memory(format!("episodic index ro_txn failed: {e}")))?;
349 let mut cursor = tx
350 .open_ro_cursor(self.term_db)
351 .map_err(|e| CoreError::Memory(format!("episodic index cursor failed: {e}")))?;
352 Ok(cursor.iter().next().is_none())
353 }
354
355 pub fn record_count(&self) -> Result<u64> {
357 let tx = self
358 .env
359 .begin_ro_txn()
360 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
361 let mut cursor = tx
362 .open_ro_cursor(self.db)
363 .map_err(|e| CoreError::Memory(format!("episodic cursor failed: {e}")))?;
364 let mut count = 0u64;
365 for _ in cursor.iter() {
366 count += 1;
367 }
368 Ok(count)
369 }
370
371 pub fn append_explicit(
373 &self,
374 record: &EpisodicRecord,
375 policy: EpisodicCapturePolicy,
376 ) -> Result<bool> {
377 let prepared = record
378 .clone()
379 .with_content(policy.prepare_content(&record.content));
380 self.append(&prepared)?;
381 Ok(true)
382 }
383
384 pub fn append_explicit_batch(
386 &self,
387 records: &[EpisodicRecord],
388 policy: EpisodicCapturePolicy,
389 ) -> Result<usize> {
390 if records.is_empty() {
391 return Ok(0);
392 }
393 let prepared: Vec<EpisodicRecord> = records
394 .iter()
395 .map(|record| {
396 record
397 .clone()
398 .with_content(policy.prepare_content(&record.content))
399 })
400 .collect();
401 self.append_batch(&prepared)?;
402 Ok(prepared.len())
403 }
404
405 pub fn get(&self, id: EpisodicId) -> Result<Option<EpisodicRecord>> {
407 let tx = self
408 .env
409 .begin_ro_txn()
410 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
411 let result = tx.get(self.db, id.as_bytes());
412 match result {
413 Ok(bytes) => {
414 let record: EpisodicRecord = rmp_serde::from_slice(bytes)
415 .map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
416 tx.commit()
417 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
418 Ok(Some(record))
419 }
420 Err(lmdb::Error::NotFound) => {
421 tx.commit()
422 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
423 Ok(None)
424 }
425 Err(e) => Err(CoreError::Memory(format!("episodic get failed: {e}"))),
426 }
427 }
428
429 pub fn replace_content(&self, id: EpisodicId, new_content: &str) -> Result<bool> {
438 let mut tx = self
439 .env
440 .begin_rw_txn()
441 .map_err(|e| CoreError::Memory(format!("episodic rw_txn failed: {e}")))?;
442 let existing = match tx.get(self.db, id.as_bytes()) {
443 Ok(bytes) => bytes.to_vec(),
444 Err(lmdb::Error::NotFound) => {
445 tx.abort();
446 return Ok(false);
447 }
448 Err(e) => {
449 tx.abort();
450 return Err(CoreError::Memory(format!("episodic get failed: {e}")));
451 }
452 };
453 let mut record: EpisodicRecord = rmp_serde::from_slice(&existing)
454 .map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
455 if record.content == new_content {
456 tx.abort();
457 return Ok(true);
458 }
459 let old_terms = self.record_terms(&record);
460 record.content = new_content.to_string();
461 let new_terms = self.record_terms(&record);
462 let value = rmp_serde::to_vec(&record)
463 .map_err(|e| CoreError::Memory(format!("episodic serialize failed: {e}")))?;
464 tx.put(self.db, id.as_bytes(), &value, WriteFlags::default())
465 .map_err(|e| CoreError::Memory(format!("episodic replace write failed: {e}")))?;
466 for term in old_terms.difference(&new_terms) {
467 match tx.del(self.term_db, term, Some(id.as_bytes())) {
468 Ok(()) | Err(lmdb::Error::NotFound) => {}
469 Err(e) => {
470 tx.abort();
471 return Err(CoreError::Memory(format!(
472 "episodic term delete failed: {e}"
473 )));
474 }
475 }
476 }
477 for term in &new_terms {
478 if let Err(e) = tx.put(self.term_db, term, id.as_bytes(), WriteFlags::default()) {
479 tx.abort();
480 return Err(CoreError::Memory(format!(
481 "episodic term write failed: {e}"
482 )));
483 }
484 }
485 tx.commit()
486 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
487 self.mutation_count
488 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
489 self.clear_term_cache();
490 Ok(true)
491 }
492
493 pub fn transition(&self, id: EpisodicId, transition: MemoryTransition) -> Result<()> {
495 let mut tx = self
496 .env
497 .begin_rw_txn()
498 .map_err(|e| CoreError::Memory(format!("episodic rw_txn failed: {e}")))?;
499 let bytes = match tx.get(self.db, id.as_bytes()) {
500 Ok(bytes) => bytes,
501 Err(lmdb::Error::NotFound) => {
502 tx.abort();
503 return Err(CoreError::InvalidArgs(format!(
504 "episodic record {id} does not exist"
505 )));
506 }
507 Err(e) => {
508 tx.abort();
509 return Err(CoreError::Memory(format!("episodic get failed: {e}")));
510 }
511 };
512 let mut record: EpisodicRecord = rmp_serde::from_slice(bytes)
513 .map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
514 record
515 .transition(transition)
516 .map_err(|e| CoreError::InvalidArgs(format!("episodic transition rejected: {e}")))?;
517 let value = rmp_serde::to_vec(&record)
518 .map_err(|e| CoreError::Memory(format!("episodic serialize failed: {e}")))?;
519 tx.put(self.db, id.as_bytes(), &value, WriteFlags::default())
520 .map_err(|e| CoreError::Memory(format!("episodic transition write failed: {e}")))?;
521 tx.commit()
522 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
523 self.mutation_count
524 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
525 Ok(())
526 }
527
528 pub fn scan(
530 &self,
531 session_id: Option<uuid::Uuid>,
532 limit: usize,
533 ) -> Result<Vec<EpisodicRecord>> {
534 if limit == 0 {
535 return Ok(Vec::new());
536 }
537 let tx = self
538 .env
539 .begin_ro_txn()
540 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
541 let mut cursor = tx
542 .open_ro_cursor(self.db)
543 .map_err(|e| CoreError::Memory(format!("episodic cursor failed: {e}")))?;
544 let mut records = Vec::new();
545 for item in cursor.iter() {
546 let (_, bytes) = item;
547 let record: EpisodicRecord = rmp_serde::from_slice(bytes)
548 .map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
549 if session_id.is_none_or(|id| record.session_id == Some(id)) {
550 records.push(record);
551 }
552 }
553 drop(cursor);
554 tx.commit()
555 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
556 records.sort_by_key(|record| (record.sequence, record.created_at, record.id));
557 records.truncate(limit);
558 Ok(records)
559 }
560
561 pub fn search(
571 &self,
572 query: &str,
573 limit: usize,
574 include_historical: bool,
575 ) -> Result<Vec<EpisodicSearchResult>> {
576 self.search_with_limits(query, limit, limit.saturating_mul(2), include_historical)
577 }
578
579 pub fn search_with_limits(
587 &self,
588 query: &str,
589 limit: usize,
590 candidate_limit: usize,
591 include_historical: bool,
592 ) -> Result<Vec<EpisodicSearchResult>> {
593 if limit == 0 {
594 return Ok(Vec::new());
595 }
596 let mut results = self.search_scored(query, limit, candidate_limit, include_historical)?;
597 if is_current_query(query) {
598 self.resolve_current(&mut results);
599 }
600 results.truncate(limit);
601 Ok(results)
602 }
603
604 fn search_scored(
609 &self,
610 query: &str,
611 limit: usize,
612 candidate_limit: usize,
613 include_historical: bool,
614 ) -> Result<Vec<EpisodicSearchResult>> {
615 let plan = QueryPlan::plan(query, limit);
616 let candidate_limit = candidate_limit.max(plan.candidate_limit);
617 let query_terms = tokenize(query);
618 let query_keys = key_index_terms_with_aliases(query, self.aliases.as_ref());
619 if query_terms.is_empty() && query_keys.is_empty() || limit == 0 {
620 return Ok(Vec::new());
621 }
622 let mut candidate_scores: HashMap<EpisodicId, usize> = HashMap::new();
623 for term in query_terms.iter().chain(query_keys.iter()) {
624 for id in self.term_postings(term)? {
625 *candidate_scores.entry(id).or_default() += 1;
626 }
627 }
628
629 let records = if candidate_scores.is_empty() {
630 if !matches!(self.sidecar_is_empty(), Ok(true)) {
635 return Ok(Vec::new());
636 }
637 self.scan(None, usize::MAX)?
639 } else {
640 let mut ranked_candidates: Vec<(EpisodicId, usize)> =
641 candidate_scores.into_iter().collect();
642 ranked_candidates.sort_by(|(left_id, left_count), (right_id, right_count)| {
643 right_count
644 .cmp(left_count)
645 .then_with(|| left_id.cmp(right_id))
646 });
647 ranked_candidates.truncate(candidate_limit);
648 self.load_records(
649 &ranked_candidates
650 .into_iter()
651 .map(|(id, _)| id)
652 .collect::<Vec<_>>(),
653 )?
654 };
655
656 let mut results = Vec::new();
657 for record in records {
658 if !include_historical && !matches!(record.validity, ValidityState::Active) {
659 continue;
660 }
661 let content_terms = tokenize(&record.content);
662 let content_keys = key_index_terms_with_aliases(&record.content, self.aliases.as_ref());
663 let reverse_map: HashMap<&String, Vec<String>> =
668 if let Some(ref enrichment) = self.enrichment {
669 if matches!(record.kind, EpisodicKind::UserStatement) {
670 query_terms
671 .iter()
672 .map(|qt| (qt, enrichment.reverse_enrich(qt)))
673 .collect()
674 } else {
675 HashMap::new()
676 }
677 } else {
678 HashMap::new()
679 };
680 let mut reverse_match_count = 0usize;
681 let matched_terms = query_terms
682 .iter()
683 .filter(|term| {
684 if content_terms.iter().any(|candidate| candidate == *term)
685 || content_keys.iter().any(|candidate| candidate == *term)
686 {
687 return true;
688 }
689 if let Some(reverse_terms) = reverse_map.get(term) {
692 let found = reverse_terms.iter().any(|rt| {
693 content_terms.iter().any(|candidate| candidate == rt)
694 || content_keys.iter().any(|candidate| candidate == rt)
695 });
696 if found {
697 reverse_match_count += 1;
698 }
699 return found;
700 }
701 false
702 })
703 .count();
704 let matched_keys = query_keys
705 .iter()
706 .filter(|term| {
707 content_keys.iter().any(|candidate| candidate == *term)
708 || content_terms.iter().any(|candidate| candidate == *term)
709 })
710 .count();
711 if matched_terms == 0 && matched_keys == 0 {
712 continue;
713 }
714 let key_bonus = if query_keys.is_empty() {
715 0.0
716 } else {
717 matched_keys as f32 / query_keys.len() as f32 * plan.key_weight
718 };
719 let role_boost = match record.kind {
720 EpisodicKind::UserStatement => 0.12,
721 _ => 0.0,
722 };
723 let effective_matched = if matches!(record.kind, EpisodicKind::UserStatement) {
724 (matched_terms + 2).min(query_terms.len())
725 } else {
726 matched_terms
727 };
728 let coverage = if query_terms.is_empty() {
729 0.0
730 } else {
731 effective_matched as f32 / query_terms.len() as f32
732 };
733 let number_bonus = if plan.number_query {
734 let has_digit = content_terms
735 .iter()
736 .any(|term| term.chars().any(|c| c.is_ascii_digit()));
737 if has_digit || contains_number_word(&record.content) {
738 0.03
739 } else {
740 0.0
741 }
742 } else {
743 0.0
744 };
745 let density = matched_terms as f32 / content_terms.len().max(1) as f32;
746 results.push(EpisodicSearchResult {
747 record,
748 score: coverage
749 + key_bonus
750 + role_boost
751 + number_bonus
752 + (reverse_match_count as f32).mul_add(0.05, density * 0.03),
753 matched_terms: matched_terms.max(matched_keys),
754 });
755 }
756 let mut session_counts: HashMap<Option<uuid::Uuid>, usize> = HashMap::new();
760 for r in &results {
761 *session_counts.entry(r.record.session_id).or_default() += 1;
762 }
763 for r in &mut results {
764 let count = session_counts
765 .get(&r.record.session_id)
766 .copied()
767 .unwrap_or(1);
768 if count > 1 {
769 r.score = 0.02f32.mul_add((count - 1).min(3) as f32, r.score);
770 }
771 }
772 let mut hash_counts: HashMap<&str, usize> = HashMap::new();
778 for r in &results {
779 *hash_counts
780 .entry(r.record.content_hash.as_str())
781 .or_default() += 1;
782 }
783 let hash_boosts: HashMap<String, f32> = results
784 .iter()
785 .map(|r| {
786 let count = hash_counts
787 .get(r.record.content_hash.as_str())
788 .copied()
789 .unwrap_or(1);
790 let boost = if count > 1 {
791 0.03 * (count - 1).min(3) as f32
792 } else {
793 0.0
794 };
795 (r.record.id.to_string(), boost)
796 })
797 .collect();
798 for r in &mut results {
799 if let Some(boost) = hash_boosts.get(&r.record.id.to_string()) {
800 r.score += boost;
801 }
802 }
803 results.sort_by(|a, b| {
804 b.score
805 .partial_cmp(&a.score)
806 .unwrap_or(std::cmp::Ordering::Equal)
807 .then_with(|| b.matched_terms.cmp(&a.matched_terms))
808 .then_with(|| a.record.content.len().cmp(&b.record.content.len()))
809 .then_with(|| a.record.sequence.cmp(&b.record.sequence))
810 .then_with(|| a.record.id.cmp(&b.record.id))
811 });
812 Ok(results)
813 }
814
815 pub fn search_with_rerank(
832 &self,
833 query: &str,
834 limit: usize,
835 candidate_limit: usize,
836 include_historical: bool,
837 alpha: f32,
838 rerank_pool: usize,
839 ) -> Result<Vec<EpisodicSearchResult>> {
840 let Some(ref embedder) = self.embedder else {
841 return self.search_with_limits(query, limit, candidate_limit, include_historical);
842 };
843 if !embedder.is_available() || limit == 0 {
844 return self.search_with_limits(query, limit, candidate_limit, include_historical);
845 }
846
847 let rerank_pool = if rerank_pool == 0 {
855 limit.max(candidate_limit).min(50)
856 } else {
857 rerank_pool.max(limit).min(50)
858 };
859 let deterministic: Vec<EpisodicSearchResult> = self
862 .search_scored(query, rerank_pool, candidate_limit, include_historical)?
863 .into_iter()
864 .take(rerank_pool)
865 .collect();
866 if deterministic.is_empty() {
867 return Ok(Vec::new());
868 }
869
870 let contents: Vec<&str> = std::iter::once(query)
872 .chain(deterministic.iter().map(|r| r.record.content.as_str()))
873 .collect();
874 let embeddings = embedder.embed_batch(&contents)?;
875 if embeddings.len() != deterministic.len() + 1 {
876 return Err(CoreError::Memory(format!(
877 "embedder returned {} vectors, expected {}",
878 embeddings.len(),
879 deterministic.len() + 1
880 )));
881 }
882 let query_vec = &embeddings[0];
883 let candidate_vecs = &embeddings[1..];
884
885 if alpha >= 2.0 {
886 let protected: Vec<EpisodicSearchResult> =
891 deterministic.into_iter().take(limit).collect();
892 let cosines: Vec<f32> = protected
893 .iter()
894 .enumerate()
895 .map(|(i, _)| cosine_sim(query_vec, &candidate_vecs[i]))
896 .collect();
897 let mut order: Vec<usize> = (0..protected.len()).collect();
898 order.sort_by(|&a, &b| {
899 cosines[b]
900 .partial_cmp(&cosines[a])
901 .unwrap_or(std::cmp::Ordering::Equal)
902 .then_with(|| a.cmp(&b))
904 });
905 let mut slots: Vec<Option<EpisodicSearchResult>> =
906 protected.into_iter().map(Some).collect();
907 let reranked: Vec<EpisodicSearchResult> =
908 order.into_iter().filter_map(|i| slots[i].take()).collect();
909 Ok(reranked)
910 } else if alpha >= 1.0 {
911 let delta = 0.05;
913 let mut reranked = deterministic;
914 let cosines: Vec<f32> = candidate_vecs
915 .iter()
916 .map(|v| cosine_sim(query_vec, v))
917 .collect();
918 let n = reranked.len();
920 for _ in 0..n {
921 let mut swapped = false;
922 for i in 0..n.saturating_sub(1) {
923 let det_gap = (reranked[i].score - reranked[i + 1].score).abs();
924 if det_gap < delta && cosines[i + 1] > cosines[i] {
925 reranked.swap(i, i + 1);
926 swapped = true;
927 }
928 }
929 if !swapped {
930 break;
931 }
932 }
933 if is_current_query(query) {
934 self.resolve_current(&mut reranked);
935 }
936 reranked.truncate(limit);
937 Ok(reranked)
938 } else {
939 let max_det = deterministic
941 .iter()
942 .map(|r| r.score)
943 .fold(0.0f32, f32::max)
944 .max(1e-9);
945
946 let mut reranked: Vec<EpisodicSearchResult> = deterministic
947 .into_iter()
948 .enumerate()
949 .map(|(i, mut r)| {
950 let cosine = cosine_sim(query_vec, &candidate_vecs[i]);
951 let det_norm = r.score / max_det;
952 r.score = alpha.mul_add(det_norm, (1.0 - alpha) * cosine);
953 r
954 })
955 .collect();
956
957 reranked.sort_by(|a, b| {
958 b.score
959 .partial_cmp(&a.score)
960 .unwrap_or(std::cmp::Ordering::Equal)
961 .then_with(|| b.matched_terms.cmp(&a.matched_terms))
962 .then_with(|| a.record.content.len().cmp(&b.record.content.len()))
963 .then_with(|| a.record.sequence.cmp(&b.record.sequence))
964 .then_with(|| a.record.id.cmp(&b.record.id))
965 });
966 if is_current_query(query) {
967 self.resolve_current(&mut reranked);
968 }
969 reranked.truncate(limit);
970 Ok(reranked)
971 }
972 }
973
974 fn resolve_current(&self, results: &mut Vec<EpisodicSearchResult>) {
1001 if results.len() < 2 {
1002 return;
1003 }
1004 let mut anchors: Vec<EpisodicSearchResult> = Vec::new();
1005 let mut rest: Vec<EpisodicSearchResult> = Vec::new();
1006 for result in results.drain(..) {
1007 let is_anchor = matches!(result.record.kind, EpisodicKind::UserStatement)
1008 && contains_change_marker(&result.record.content);
1009 if is_anchor {
1010 anchors.push(result);
1011 } else {
1012 rest.push(result);
1013 }
1014 }
1015 if anchors.is_empty() {
1016 *results = rest;
1018 return;
1019 }
1020 anchors.sort_by(|a, b| {
1021 b.record
1022 .created_at
1023 .cmp(&a.record.created_at)
1024 .then_with(|| b.record.sequence.cmp(&a.record.sequence))
1025 .then_with(|| a.record.id.cmp(&b.record.id))
1026 });
1027 anchors.extend(rest);
1028 *results = anchors;
1029 }
1030
1031 fn load_records(&self, ids: &[EpisodicId]) -> Result<Vec<EpisodicRecord>> {
1032 let tx = self
1033 .env
1034 .begin_ro_txn()
1035 .map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
1036 let mut records = Vec::with_capacity(ids.len());
1037 for id in ids {
1038 match tx.get(self.db, id.as_bytes()) {
1039 Ok(bytes) => {
1040 records.push(rmp_serde::from_slice(bytes).map_err(|e| {
1041 CoreError::Memory(format!("episodic deserialize failed: {e}"))
1042 })?);
1043 }
1044 Err(lmdb::Error::NotFound) => {}
1045 Err(e) => return Err(CoreError::Memory(format!("episodic get failed: {e}"))),
1046 }
1047 }
1048 tx.commit()
1049 .map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
1050 Ok(records)
1051 }
1052}
1053
1054fn index_terms_with_aliases(text: &str, aliases: Option<&AdaptiveAliases>) -> Vec<String> {
1055 tokenize(text)
1056 .into_iter()
1057 .chain(key_index_terms_with_aliases(text, aliases))
1058 .fold(Vec::new(), |mut terms, term| {
1059 if !terms.contains(&term) {
1060 terms.push(term);
1061 }
1062 terms
1063 })
1064}
1065
1066const MAX_TERM_KEY_BYTES: usize = 480;
1076
1077fn index_safe_term(term: &str) -> String {
1082 if term.len() <= MAX_TERM_KEY_BYTES {
1083 return term.to_string();
1084 }
1085 let digest = Sha256::digest(term.as_bytes());
1086 format!("~h:{digest:x}")
1087}
1088
1089fn tokenize(text: &str) -> Vec<String> {
1090 strip_stopwords(text)
1091 .split(|c: char| !c.is_alphanumeric())
1092 .filter(|term| term.len() > 1)
1093 .map(|term| simple_stem(&term.to_ascii_lowercase()))
1094 .fold(Vec::new(), |mut terms, term| {
1095 if !terms.contains(&term) {
1096 terms.push(term);
1097 }
1098 terms
1099 })
1100}
1101
1102const CURRENT_QUERY_WORD_CUES: &[&str] = &["current", "currently", "latest", "nowadays"];
1104
1105const CURRENT_QUERY_PHRASE_CUES: &[&str] = &["these days", "right now", "at the moment"];
1107
1108#[must_use]
1115pub fn is_current_query(query: &str) -> bool {
1116 let lowered = query.to_ascii_lowercase();
1117 let has_word = lowered
1118 .split(|c: char| !c.is_alphanumeric())
1119 .any(|token| CURRENT_QUERY_WORD_CUES.contains(&token));
1120 has_word
1121 || CURRENT_QUERY_PHRASE_CUES
1122 .iter()
1123 .any(|cue| lowered.contains(cue))
1124}
1125
1126const CHANGE_MARKERS: &[&str] = &[
1131 "switched to",
1132 "switch to",
1133 "switching to",
1134 "switched from",
1135 "changed my",
1136 "change my",
1137 "changed from",
1138 "now prefer",
1139 "now i prefer",
1140 "now i'm",
1141 "now im",
1142 "no longer",
1143 "used to",
1144 "moved to",
1145 "not anymore",
1146 "instead of",
1147 "replaced",
1148 "gave up",
1149];
1150
1151fn contains_change_marker(content: &str) -> bool {
1153 let lowered = content.to_ascii_lowercase();
1154 CHANGE_MARKERS.iter().any(|marker| lowered.contains(marker))
1155}
1156
1157const CONTRADICTION_MARKERS: &[&str] = &[
1161 "no longer",
1162 "anymore",
1163 "changed my mind",
1164 "changed my",
1165 "used to",
1166 "gave up",
1167 "just a phase",
1168 "not really",
1169 "but i",
1170];
1171
1172#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1179pub struct EpisodicConflict {
1180 pub later_record: EpisodicId,
1183 pub earlier_record: EpisodicId,
1185 pub marker: String,
1187 pub shared_terms: Vec<String>,
1189 pub later_content: String,
1191 pub earlier_content: String,
1193}
1194
1195#[must_use]
1203pub fn detect_conflicts(results: &[EpisodicSearchResult]) -> Vec<EpisodicConflict> {
1204 const MAX_CONFLICTS: usize = 10;
1205 let mut conflicts: Vec<EpisodicConflict> = Vec::new();
1206 let mut seen_pairs: Vec<(EpisodicId, EpisodicId)> = Vec::new();
1207 for (i, marked) in results.iter().enumerate() {
1208 if !matches!(marked.record.kind, EpisodicKind::UserStatement) {
1209 continue;
1210 }
1211 let lowered = marked.record.content.to_ascii_lowercase();
1212 let Some(marker) = CONTRADICTION_MARKERS
1213 .iter()
1214 .find(|m| lowered.contains(*m))
1215 .copied()
1216 else {
1217 continue;
1218 };
1219 let marked_terms = tokenize(&marked.record.content);
1220 for (j, other) in results.iter().enumerate() {
1221 if i == j || !matches!(other.record.kind, EpisodicKind::UserStatement) {
1222 continue;
1223 }
1224 if other.record.content_hash == marked.record.content_hash {
1225 continue;
1226 }
1227 let other_terms = tokenize(&other.record.content);
1228 let shared: Vec<String> = marked_terms
1229 .iter()
1230 .filter(|t| other_terms.contains(t))
1231 .cloned()
1232 .collect();
1233 if shared.len() < 2 {
1234 continue;
1235 }
1236 let (later, earlier) = if (marked.record.created_at, marked.record.sequence)
1239 > (other.record.created_at, other.record.sequence)
1240 {
1241 (&marked.record, &other.record)
1242 } else {
1243 (&other.record, &marked.record)
1244 };
1245 let pair_key = if later.id < earlier.id {
1246 (later.id, earlier.id)
1247 } else {
1248 (earlier.id, later.id)
1249 };
1250 if seen_pairs.contains(&pair_key) {
1251 continue;
1252 }
1253 seen_pairs.push(pair_key);
1254 conflicts.push(EpisodicConflict {
1255 later_record: later.id,
1256 earlier_record: earlier.id,
1257 marker: marker.to_string(),
1258 shared_terms: shared,
1259 later_content: later.content.clone(),
1260 earlier_content: earlier.content.clone(),
1261 });
1262 if conflicts.len() >= MAX_CONFLICTS {
1263 return conflicts;
1264 }
1265 }
1266 }
1267 conflicts
1268}
1269
1270fn simple_stem(word: &str) -> String {
1271 if word.len() <= 3 {
1272 return word.to_string();
1273 }
1274 for suffix in ["ies", "ied", "ing", "edly", "ed", "ly", "es", "s"] {
1275 if let Some(stem) = word.strip_suffix(suffix) {
1276 if suffix == "ies" || suffix == "ied" {
1277 return format!("{stem}y");
1278 }
1279 if stem.len() >= 2 {
1280 return stem.to_string();
1281 }
1282 }
1283 }
1284 word.to_string()
1285}
1286
1287fn cosine_sim(a: &[f32], b: &[f32]) -> f32 {
1288 let dot = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum::<f32>();
1289 let norm_a = a.iter().map(|x| x * x).sum::<f32>().sqrt();
1290 let norm_b = b.iter().map(|x| x * x).sum::<f32>().sqrt();
1291 if norm_a < 1e-9 || norm_b < 1e-9 {
1292 0.0
1293 } else {
1294 dot / (norm_a * norm_b)
1295 }
1296}
1297
1298fn contains_number_word(text: &str) -> bool {
1299 const NUMBER_WORDS: &[&str] = &[
1300 "one",
1301 "two",
1302 "three",
1303 "four",
1304 "five",
1305 "six",
1306 "seven",
1307 "eight",
1308 "nine",
1309 "ten",
1310 "eleven",
1311 "twelve",
1312 "thirteen",
1313 "fourteen",
1314 "fifteen",
1315 "sixteen",
1316 "seventeen",
1317 "eighteen",
1318 "nineteen",
1319 "twenty",
1320 "thirty",
1321 "forty",
1322 "fifty",
1323 "sixty",
1324 "seventy",
1325 "eighty",
1326 "ninety",
1327 "hundred",
1328 "thousand",
1329 "million",
1330 "billion",
1331 "dozen",
1332 "couple",
1333 "half",
1334 "quarter",
1335 "double",
1336 "triple",
1337 "twice",
1338 ];
1339 for word in text.split(|c: char| !c.is_alphanumeric()) {
1340 if word.len() >= 3 && NUMBER_WORDS.iter().any(|nw| word.eq_ignore_ascii_case(nw)) {
1341 return true;
1342 }
1343 }
1344 false
1345}
1346
1347#[cfg(test)]
1348mod tests {
1349 use super::*;
1350 use crate::MemoryStore;
1351 #[cfg(unix)]
1355 use chrono::{DateTime, Utc};
1356 use tempfile::tempdir;
1357 #[cfg(unix)]
1358 use uuid::Uuid;
1359 use wm_core::{EpisodicKind, Provenance, ProvenanceSource, ValidityState};
1360
1361 fn sample_record(sequence: u64, content: &str) -> EpisodicRecord {
1362 EpisodicRecord::new(
1363 None,
1364 sequence,
1365 EpisodicKind::Observation,
1366 content,
1367 Provenance::new(ProvenanceSource::User),
1368 )
1369 }
1370
1371 fn user_statement(sequence: u64, content: &str) -> EpisodicRecord {
1372 EpisodicRecord::new(
1373 None,
1374 sequence,
1375 EpisodicKind::UserStatement,
1376 content,
1377 Provenance::new(ProvenanceSource::User),
1378 )
1379 }
1380
1381 fn assistant_response(sequence: u64, content: &str) -> EpisodicRecord {
1382 EpisodicRecord::new(
1383 None,
1384 sequence,
1385 EpisodicKind::AssistantResponse,
1386 content,
1387 Provenance::new(ProvenanceSource::Agent),
1388 )
1389 }
1390
1391 #[cfg(unix)]
1399 fn q06_record(id: u128, sequence: u64, content: &str, created_at: &str) -> EpisodicRecord {
1400 let mut record = EpisodicRecord::new(
1401 None,
1402 sequence,
1403 EpisodicKind::Observation,
1404 content,
1405 Provenance::new(ProvenanceSource::User),
1406 )
1407 .with_id(Uuid::from_u128(id));
1408 record.created_at = DateTime::parse_from_rfc3339(created_at)
1409 .unwrap()
1410 .with_timezone(&Utc);
1411 record
1412 }
1413
1414 #[cfg(unix)]
1415 fn q06_acknowledged() -> EpisodicRecord {
1416 q06_record(601, 601, "q06 acknowledged control", "2026-01-01T00:10:01Z")
1417 }
1418
1419 #[cfg(unix)]
1420 fn run_q06_child(case: &str, store_path: &std::path::Path) {
1421 use std::io::Write;
1422 let expected_uuid = std::env::var("WM_Q06_UUID").expect("WM_Q06_UUID");
1423 let uuid = Uuid::parse_str(&expected_uuid).expect("WM_Q06_UUID must parse");
1424 let (sequence, content, created_at) = match case {
1425 "before_raw_commit" => (602, "q06 precommit candidate", "2026-01-01T00:10:02Z"),
1426 "after_raw_commit" => (603, "q06 uncertain candidate", "2026-01-01T00:10:03Z"),
1427 other => panic!("q06 child: unknown case {other}"),
1428 };
1429 let record = q06_record(uuid.as_u128(), sequence, content, created_at);
1430 let store = MemoryStore::open_default(store_path).expect("q06 child store");
1431 match store.episodic().append(&record) {
1432 Ok(()) => {
1433 println!("WM_Q06_CALLER_ACK {uuid}");
1436 let _ = std::io::stdout().flush();
1437 }
1438 Err(e) => {
1439 eprintln!("q06 child append failed: {e}");
1440 std::process::exit(2);
1441 }
1442 }
1443 }
1444
1445 #[cfg(unix)]
1446 fn run_q06_killed_case(
1447 test_filter: &str,
1448 store_path: &std::path::Path,
1449 case: &str,
1450 uuid: Uuid,
1451 ) -> Vec<String> {
1452 use std::io::{BufRead, BufReader};
1453 let exe = std::env::current_exe().expect("q06 current_exe");
1454 let mut child = std::process::Command::new(exe)
1455 .arg(test_filter)
1456 .arg("--exact")
1457 .arg("--nocapture")
1458 .env("WM_Q06_CASE", case)
1459 .env("WM_Q06_STORE", store_path)
1460 .env("WM_Q06_UUID", uuid.to_string())
1461 .stdout(std::process::Stdio::piped())
1462 .stdin(std::process::Stdio::piped())
1463 .stderr(std::process::Stdio::inherit())
1464 .spawn()
1465 .expect("q06 spawn child");
1466
1467 let stdout = child.stdout.take().expect("q06 child stdout");
1468 let (tx, rx) = std::sync::mpsc::channel::<String>();
1469 let reader = std::thread::spawn(move || {
1470 for line in BufReader::new(stdout).lines() {
1471 match line {
1472 Ok(line) => {
1473 if tx.send(line).is_err() {
1474 break;
1475 }
1476 }
1477 Err(_) => break,
1478 }
1479 }
1480 });
1481
1482 let expected = format!("WM_Q06_BOUNDARY {case} {uuid}");
1483 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
1484 let mut lines = Vec::new();
1485 let mut seen = false;
1486 while !seen {
1487 let remaining = deadline.saturating_duration_since(std::time::Instant::now());
1488 if remaining.is_zero() {
1489 break;
1490 }
1491 match rx.recv_timeout(remaining) {
1492 Ok(line) if line == expected => seen = true,
1493 Ok(line) => lines.push(line),
1494 Err(_) => break,
1495 }
1496 }
1497 if !seen {
1498 let _ = child.kill();
1499 let _ = child.wait();
1500 reader.join().ok();
1501 panic!("q06 case {case}: boundary {expected:?} not observed; lines={lines:?}");
1502 }
1503
1504 child.kill().expect("q06 kill blocked child");
1505 let status = child.wait().expect("q06 reap child");
1506 assert!(
1507 !status.success(),
1508 "q06 case {case}: killed child must not exit successfully: {status:?}"
1509 );
1510 reader.join().ok();
1511 while let Ok(line) = rx.try_recv() {
1512 lines.push(line);
1513 }
1514 assert!(
1515 !lines.iter().any(|line| line.contains("WM_Q06_CALLER_ACK")),
1516 "q06 case {case}: caller acknowledgement in a killed case invalidates the experiment: {lines:?}"
1517 );
1518 lines
1519 }
1520
1521 #[cfg(unix)]
1522 #[test]
1523 fn q06_commit_boundary_sigkill_classification() {
1524 if let Ok(case) = std::env::var("WM_Q06_CASE") {
1525 let store_path = std::env::var("WM_Q06_STORE").expect("WM_Q06_STORE");
1526 run_q06_child(&case, std::path::Path::new(&store_path));
1527 return;
1528 }
1529
1530 let dir = tempdir().unwrap();
1531 let store_path = dir.path().join("lmdb");
1532 let acknowledged = q06_acknowledged();
1533
1534 {
1536 let store = MemoryStore::open_default(&store_path).unwrap();
1537 store.episodic().append(&acknowledged).unwrap();
1538 let read = store
1539 .episodic()
1540 .get(acknowledged.id)
1541 .unwrap()
1542 .expect("acknowledged record");
1543 assert_eq!(read, acknowledged);
1544 }
1545
1546 let test_filter = "episodic::tests::q06_commit_boundary_sigkill_classification";
1547
1548 let candidate_before =
1550 q06_record(602, 602, "q06 precommit candidate", "2026-01-01T00:10:02Z");
1551 run_q06_killed_case(
1552 test_filter,
1553 &store_path,
1554 "before_raw_commit",
1555 candidate_before.id,
1556 );
1557 {
1558 let store = MemoryStore::open_default(&store_path).unwrap();
1559 assert_eq!(
1560 store.episodic().get(acknowledged.id).unwrap(),
1561 Some(acknowledged.clone()),
1562 "acknowledged record must survive a pre-commit kill unchanged"
1563 );
1564 assert_eq!(
1565 store.episodic().get(candidate_before.id).unwrap(),
1566 None,
1567 "pre-commit candidate must be absent after reopen"
1568 );
1569 }
1570
1571 let candidate_after =
1574 q06_record(603, 603, "q06 uncertain candidate", "2026-01-01T00:10:03Z");
1575 run_q06_killed_case(
1576 test_filter,
1577 &store_path,
1578 "after_raw_commit",
1579 candidate_after.id,
1580 );
1581 {
1582 let store = MemoryStore::open_default(&store_path).unwrap();
1583 assert_eq!(
1584 store.episodic().get(acknowledged.id).unwrap(),
1585 Some(acknowledged),
1586 "acknowledged record must survive a post-commit kill unchanged"
1587 );
1588 assert_eq!(
1589 store.episodic().get(candidate_before.id).unwrap(),
1590 None,
1591 "pre-commit candidate must stay absent"
1592 );
1593 assert_eq!(
1594 store.episodic().get(candidate_after.id).unwrap(),
1595 Some(candidate_after),
1596 "post-commit candidate must be present and byte-equal after reopen"
1597 );
1598 }
1599 }
1600
1601 #[test]
1602 fn current_query_detection() {
1603 assert!(is_current_query("What's my current favorite coffee?"));
1604 assert!(is_current_query("What am I currently reading these days?"));
1605 assert!(is_current_query("What's the latest book I mentioned?"));
1606 assert!(is_current_query("What's my job right now?"));
1607 assert!(is_current_query("What am I eating at the moment?"));
1608 assert!(!is_current_query("What's my favorite coffee?"));
1609 assert!(!is_current_query("Where did I volunteer in February?"));
1610 assert!(!is_current_query("What did I say about the trip?"));
1611 assert!(!is_current_query("What currency did I use in Japan?"));
1613 }
1614
1615 #[test]
1616 fn protected_rerank_preserves_membership_with_a_small_pool() {
1617 use crate::embedder::Embedder;
1622 use std::sync::Arc;
1623 use std::sync::atomic::{AtomicUsize, Ordering};
1624
1625 struct CountingEmbedder(Arc<AtomicUsize>);
1626 impl Embedder for CountingEmbedder {
1627 fn embed_batch(&self, texts: &[&str]) -> wm_core::Result<Vec<Vec<f32>>> {
1628 self.0.store(texts.len(), Ordering::Relaxed);
1629 Ok(texts
1630 .iter()
1631 .map(|t| {
1632 let mut v = vec![0.0_f32; 8];
1633 for (i, b) in t.bytes().enumerate() {
1634 v[i % 8] += f32::from(b) / 255.0;
1635 }
1636 v
1637 })
1638 .collect())
1639 }
1640 fn dimension(&self) -> usize {
1641 8
1642 }
1643 fn is_available(&self) -> bool {
1644 true
1645 }
1646 fn backend_name(&self) -> &'static str {
1647 "counting-test"
1648 }
1649 }
1650
1651 let tmp = tempdir().unwrap();
1652 let store = MemoryStore::open_default(tmp.path()).unwrap();
1653 for i in 0..20 {
1654 store
1655 .episodic()
1656 .append(&user_statement(
1657 i,
1658 &format!("Record {i} discusses topic alpha beta gamma delta epsilon"),
1659 ))
1660 .unwrap();
1661 }
1662
1663 let query = "topic alpha beta gamma delta epsilon";
1664 let expected: Vec<uuid::Uuid> = store
1667 .episodic()
1668 .search_scored(query, 5, 20, false)
1669 .unwrap()
1670 .iter()
1671 .take(5)
1672 .map(|r| r.record.id)
1673 .collect();
1674 assert_eq!(expected.len(), 5);
1675
1676 let batches = Arc::new(AtomicUsize::new(0));
1677 store.set_episodic_embedder(Arc::new(CountingEmbedder(batches.clone())));
1678 let reranked = store
1679 .episodic()
1680 .search_with_rerank(query, 5, 20, false, 2.0, 5)
1681 .unwrap();
1682 assert_eq!(
1683 batches.load(Ordering::Relaxed),
1684 6,
1685 "small rerank pool must embed pool + query only"
1686 );
1687 let mut got: Vec<uuid::Uuid> = reranked.iter().map(|r| r.record.id).collect();
1688 let mut want = expected;
1689 got.sort();
1690 want.sort();
1691 assert_eq!(
1692 got, want,
1693 "protected mode must keep the deterministic candidate set (membership), only reorder it"
1694 );
1695 }
1696
1697 #[test]
1698 fn current_query_resolution_prefers_latest_statement() {
1699 let tmp = tempdir().unwrap();
1700 let store = MemoryStore::open_default(tmp.path()).unwrap();
1701 let episodic = store.episodic();
1702 episodic
1706 .append(&user_statement(1, "My favorite coffee is dark roast."))
1707 .unwrap();
1708 episodic
1709 .append(&user_statement(
1710 2,
1711 "I really love dark roast when it comes to coffee.",
1712 ))
1713 .unwrap();
1714 episodic
1715 .append(&user_statement(3, "I've been jogging lately."))
1716 .unwrap();
1717 episodic
1718 .append(&user_statement(4, "I've switched to cold brew for coffee."))
1719 .unwrap();
1720
1721 let results = episodic
1722 .search("What's my current favorite coffee?", 5, false)
1723 .unwrap();
1724 assert!(!results.is_empty());
1725 assert!(
1726 results[0].record.content.contains("cold brew"),
1727 "current query must rank the latest statement first, got: {}",
1728 results[0].record.content
1729 );
1730 }
1731
1732 #[test]
1733 fn current_query_anchors_switched_from_template() {
1734 let tmp = tempdir().unwrap();
1740 let store = MemoryStore::open_default(tmp.path()).unwrap();
1741 let episodic = store.episodic();
1742 episodic
1743 .append(&user_statement(1, "My favorite coffee is espresso."))
1744 .unwrap();
1745 episodic
1746 .append(&user_statement(
1747 2,
1748 "I've actually switched from espresso to latte for coffee.",
1749 ))
1750 .unwrap();
1751 episodic
1752 .append(&user_statement(3, "My favorite coffee is latte."))
1753 .unwrap();
1754
1755 let results = episodic
1756 .search("What's my current favorite coffee?", 5, false)
1757 .unwrap();
1758 assert!(!results.is_empty());
1759 assert!(
1760 results[0].record.content.contains("latte"),
1761 "'switched from' must anchor the current value, got: {}",
1762 results[0].record.content
1763 );
1764 }
1765
1766 #[test]
1767 fn non_current_query_keeps_score_order() {
1768 let tmp = tempdir().unwrap();
1769 let store = MemoryStore::open_default(tmp.path()).unwrap();
1770 let episodic = store.episodic();
1771 episodic
1772 .append(&user_statement(1, "My favorite coffee is dark roast."))
1773 .unwrap();
1774 episodic
1775 .append(&user_statement(2, "I've switched to cold brew for coffee."))
1776 .unwrap();
1777
1778 let results = episodic
1781 .search("What's my favorite coffee?", 5, false)
1782 .unwrap();
1783 assert!(!results.is_empty());
1784 assert!(
1785 results[0].record.content.contains("dark roast"),
1786 "non-current query must keep score order, got: {}",
1787 results[0].record.content
1788 );
1789 }
1790
1791 #[test]
1792 fn current_resolution_anchors_on_user_statements_only() {
1793 let tmp = tempdir().unwrap();
1794 let store = MemoryStore::open_default(tmp.path()).unwrap();
1795 let episodic = store.episodic();
1796 episodic
1797 .append(&user_statement(1, "My favorite coffee is dark roast."))
1798 .unwrap();
1799 episodic
1801 .append(&assistant_response(
1802 2,
1803 "Got it, dark roast is your favorite coffee!",
1804 ))
1805 .unwrap();
1806 episodic
1807 .append(&user_statement(3, "I've switched to cold brew for coffee."))
1808 .unwrap();
1809
1810 let results = episodic
1811 .search("What's my current favorite coffee?", 5, false)
1812 .unwrap();
1813 assert!(
1814 results[0].record.content.contains("cold brew"),
1815 "user statements anchor chronology, got: {}",
1816 results[0].record.content
1817 );
1818 assert_eq!(results[0].record.kind, EpisodicKind::UserStatement);
1819 }
1820
1821 #[test]
1822 fn current_query_without_change_markers_keeps_score_order() {
1823 let tmp = tempdir().unwrap();
1824 let store = MemoryStore::open_default(tmp.path()).unwrap();
1825 let episodic = store.episodic();
1826 episodic
1828 .append(&user_statement(
1829 1,
1830 "My favorite hiking trail is Eagle Ridge.",
1831 ))
1832 .unwrap();
1833 episodic
1834 .append(&user_statement(2, "I go hiking every weekend."))
1835 .unwrap();
1836
1837 let results = episodic
1838 .search("What's my current favorite hiking trail?", 5, false)
1839 .unwrap();
1840 assert!(!results.is_empty());
1841 assert!(
1842 results[0].record.content.contains("Eagle Ridge"),
1843 "no change markers → deterministic score order, got: {}",
1844 results[0].record.content
1845 );
1846 }
1847
1848 #[test]
1849 fn detect_conflicts_flags_contradiction_with_shared_topic() {
1850 let tmp = tempdir().unwrap();
1851 let store = MemoryStore::open_default(tmp.path()).unwrap();
1852 let episodic = store.episodic();
1853 episodic
1854 .append(&user_statement(
1855 1,
1856 "I'm vegetarian now. I decided to stop eating animal products.",
1857 ))
1858 .unwrap();
1859 episodic
1860 .append(&user_statement(
1861 2,
1862 "I'm not really vegetarian anymore, I eat steak now.",
1863 ))
1864 .unwrap();
1865 episodic
1866 .append(&user_statement(3, "I went hiking yesterday."))
1867 .unwrap();
1868
1869 let results = episodic
1870 .search("vegetarian steak eating", 10, false)
1871 .unwrap();
1872 let conflicts = detect_conflicts(&results);
1873 assert_eq!(
1874 conflicts.len(),
1875 1,
1876 "the vegetarian/steak pair must be flagged, got {conflicts:?}"
1877 );
1878 let conflict = &conflicts[0];
1879 assert!(conflict.later_content.contains("steak"));
1880 assert!(conflict.earlier_content.contains("animal products"));
1881 assert!(
1882 conflict.shared_terms.iter().any(|t| t == "vegetarian"),
1883 "shared terms must include the topic: {:?}",
1884 conflict.shared_terms
1885 );
1886 }
1887
1888 #[test]
1889 fn detect_conflicts_ignores_plain_statements_and_assistant_turns() {
1890 let tmp = tempdir().unwrap();
1891 let store = MemoryStore::open_default(tmp.path()).unwrap();
1892 let episodic = store.episodic();
1893 episodic
1895 .append(&user_statement(1, "My favorite coffee is dark roast."))
1896 .unwrap();
1897 episodic
1898 .append(&user_statement(2, "I love coffee with breakfast."))
1899 .unwrap();
1900 episodic
1902 .append(&assistant_response(
1903 3,
1904 "You mentioned you no longer like tea!",
1905 ))
1906 .unwrap();
1907
1908 let results = episodic.search("coffee tea breakfast", 10, false).unwrap();
1909 assert!(detect_conflicts(&results).is_empty());
1910 }
1911
1912 #[test]
1913 fn detect_conflicts_skips_identical_content() {
1914 let tmp = tempdir().unwrap();
1915 let store = MemoryStore::open_default(tmp.path()).unwrap();
1916 let episodic = store.episodic();
1917 let record = user_statement(1, "I'm vegetarian now, but I changed my mind.");
1918 let duplicate = user_statement(2, "I'm vegetarian now, but I changed my mind.");
1919 episodic.append(&record).unwrap();
1920 episodic.append(&duplicate).unwrap();
1921
1922 let results = episodic.search("vegetarian", 10, false).unwrap();
1923 assert!(detect_conflicts(&results).is_empty());
1925 }
1926
1927 #[test]
1928 fn append_get_transition_and_reopen_roundtrip() {
1929 let tmp = tempdir().unwrap();
1930 let session = uuid::Uuid::new_v4();
1931 let record = EpisodicRecord::new(
1932 Some(session),
1933 2,
1934 EpisodicKind::Decision,
1935 "use the raw episodic lane",
1936 Provenance::new(ProvenanceSource::User).with_actor("test"),
1937 );
1938 let id = record.id;
1939 {
1940 let store = MemoryStore::open_default(tmp.path()).unwrap();
1941 let episodic = store.episodic();
1942 episodic.append(&record).unwrap();
1943 assert_eq!(episodic.get(id).unwrap().unwrap(), record);
1944 episodic
1945 .transition(
1946 id,
1947 MemoryTransition::Supersede {
1948 replacement: uuid::Uuid::new_v4(),
1949 },
1950 )
1951 .unwrap();
1952 assert!(matches!(
1953 episodic.get(id).unwrap().unwrap().validity,
1954 ValidityState::Superseded { .. }
1955 ));
1956 }
1957 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
1958 let records = reopened.episodic().scan(Some(session), 10).unwrap();
1959 assert_eq!(records.len(), 1);
1960 assert_eq!(records[0].id, id);
1961 }
1962
1963 #[test]
1964 fn duplicate_append_is_rejected() {
1965 let tmp = tempdir().unwrap();
1966 let store = MemoryStore::open_default(tmp.path()).unwrap();
1967 let record = sample_record(1, "once");
1968 store.episodic().append(&record).unwrap();
1969 let error = store.episodic().append(&record).unwrap_err();
1970 assert!(error.to_string().contains("already exists"));
1971 }
1972
1973 #[test]
1974 fn raw_search_returns_canonical_records_and_skips_revoked_by_default() {
1975 let tmp = tempdir().unwrap();
1976 let store = MemoryStore::open_default(tmp.path()).unwrap();
1977 let active = sample_record(1, "Rust memory retrieval");
1978 let revoked = sample_record(2, "Rust memory retrieval old");
1979 let revoked_id = revoked.id;
1980 store.episodic().append(&active).unwrap();
1981 store.episodic().append(&revoked).unwrap();
1982 store
1983 .episodic()
1984 .transition(
1985 revoked_id,
1986 MemoryTransition::Revoke {
1987 reason: "stale".into(),
1988 },
1989 )
1990 .unwrap();
1991
1992 let current = store
1993 .episodic()
1994 .search("memory retrieval", 10, false)
1995 .unwrap();
1996 assert_eq!(current.len(), 1);
1997 assert_eq!(current[0].record.id, active.id);
1998
1999 let all = store
2000 .episodic()
2001 .search("memory retrieval", 10, true)
2002 .unwrap();
2003 assert_eq!(all.len(), 2);
2004 }
2005
2006 #[test]
2011 fn overlong_terms_index_and_search_without_bad_valsize() {
2012 let tmp = tempdir().unwrap();
2013 let store = MemoryStore::open_default(tmp.path()).unwrap();
2014 let long_title = "x".repeat(512);
2015 let record = sample_record(1, &format!("session title {long_title}"));
2016 store.episodic().append(&record).unwrap();
2017
2018 assert_eq!(
2021 store.episodic().get(record.id).unwrap(),
2022 Some(record.clone())
2023 );
2024
2025 let hits = store.episodic().search(&long_title, 10, false).unwrap();
2027 assert_eq!(hits.len(), 1, "long-token query must still match");
2028 assert_eq!(hits[0].record.id, record.id);
2029
2030 assert_eq!(index_safe_term(&"a".repeat(480)).len(), 480);
2032 let hashed = index_safe_term(&"a".repeat(481));
2033 assert!(hashed.starts_with("~h:"), "{hashed}");
2034 assert_eq!(hashed.len(), 67);
2035 }
2036
2037 #[test]
2041 fn sidecar_health_exposes_incomplete_derived_index_and_rebuild_heals() {
2042 let tmp = tempdir().unwrap();
2043 let store = MemoryStore::open_default(tmp.path()).unwrap();
2044 store
2045 .episodic()
2046 .append(&sample_record(1, "durable raw fact"))
2047 .unwrap();
2048 assert_eq!(store.episodic_sidecar_health().unwrap(), (1, false));
2049
2050 {
2053 let db = store.env().open_db(Some("episodic_terms_v2")).unwrap();
2054 let mut tx = store.env().begin_rw_txn().unwrap();
2055 tx.clear_db(db).unwrap();
2056 tx.commit().unwrap();
2057 }
2058 assert_eq!(
2059 store.episodic_sidecar_health().unwrap(),
2060 (1, true),
2061 "records without postings must be visible as degraded"
2062 );
2063
2064 let rebuilt = store.episodic().rebuild_sidecar().unwrap();
2066 assert_eq!(rebuilt, 1);
2067 assert_eq!(store.episodic_sidecar_health().unwrap(), (1, false));
2068 let hits = store.episodic().search("durable fact", 10, false).unwrap();
2069 assert_eq!(hits.len(), 1);
2070 }
2071
2072 #[test]
2073 fn append_batch_indexes_once_and_preserves_search() {
2074 let tmp = tempdir().unwrap();
2075 let store = MemoryStore::open_default(tmp.path()).unwrap();
2076 let first = sample_record(1, "Dr. Patel scheduled a follow-up appointment");
2077 let second = sample_record(2, "unrelated grocery list");
2078 let first_id = first.id;
2079 store.episodic().append_batch(&[first, second]).unwrap();
2080 let hits = store
2081 .episodic()
2082 .search("patel appointment", 10, false)
2083 .unwrap();
2084 assert_eq!(hits.len(), 1);
2085 assert_eq!(hits[0].record.id, first_id);
2086 }
2087
2088 #[test]
2089 fn rejected_append_batch_preserves_prior_state_after_reopen() {
2090 let tmp = tempdir().unwrap();
2094 let original = sample_record(1, "acknowledged original record");
2095 let rejected_new = sample_record(2, "must not survive rejected batch");
2096 {
2097 let store = MemoryStore::open_default(tmp.path()).unwrap();
2098 store.episodic().append(&original).unwrap();
2099 let error = store
2100 .episodic()
2101 .append_batch(&[rejected_new.clone(), original.clone()])
2102 .unwrap_err();
2103 assert!(error.to_string().contains("already exists"));
2104 assert_eq!(
2105 store.episodic().get(original.id).unwrap(),
2106 Some(original.clone())
2107 );
2108 assert_eq!(store.episodic().get(rejected_new.id).unwrap(), None);
2109 }
2110
2111 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
2112 assert_eq!(
2113 reopened.episodic().get(original.id).unwrap(),
2114 Some(original)
2115 );
2116 assert_eq!(reopened.episodic().get(rejected_new.id).unwrap(), None);
2117 }
2118
2119 #[test]
2120 fn lost_episodic_sidecar_rebuilds_from_raw_records_after_reopen() {
2121 let tmp = tempdir().unwrap();
2126 let record = sample_record(1, "sidecar recovery preserves searchable evidence");
2127 {
2128 let store = MemoryStore::open_default(tmp.path()).unwrap();
2129 let episodic = store.episodic();
2130 episodic.append(&record).unwrap();
2131 assert!(!episodic.sidecar_is_empty().unwrap());
2132
2133 let mut tx = store.env().begin_rw_txn().unwrap();
2134 tx.clear_db(episodic.term_db).unwrap();
2135 tx.commit().unwrap();
2136 assert!(episodic.sidecar_is_empty().unwrap());
2137 }
2138
2139 let reopened = MemoryStore::open_default(tmp.path()).unwrap();
2140 let episodic = reopened.episodic();
2141 assert_eq!(episodic.get(record.id).unwrap(), Some(record.clone()));
2142 assert!(!episodic.sidecar_is_empty().unwrap());
2143 let hits = episodic
2144 .search("sidecar searchable evidence", 10, false)
2145 .unwrap();
2146 assert_eq!(hits.len(), 1);
2147 assert_eq!(hits[0].record.id, record.id);
2148 }
2149
2150 #[test]
2151 fn append_explicit_batch_redacts_and_skips_private() {
2152 let tmp = tempdir().unwrap();
2153 let store = MemoryStore::open_default(tmp.path()).unwrap();
2154 let public = sample_record(1, "api_key=supersecret rust retrieval");
2155 let private = sample_record(2, "private rust retrieval").with_visibility(true, false);
2156 let public_id = public.id;
2157 store
2158 .episodic()
2159 .append_explicit_batch(&[public, private], EpisodicCapturePolicy::explicit_only())
2160 .unwrap();
2161 let stored = store.episodic().get(public_id).unwrap().unwrap();
2162 assert!(stored.content.contains("<REDACTED>"));
2163 let hits = store
2164 .episodic()
2165 .search("rust retrieval", 10, false)
2166 .unwrap();
2167 assert_eq!(hits.len(), 1);
2168 assert_eq!(hits[0].record.id, public_id);
2169 }
2170
2171 #[test]
2172 fn typed_keys_retrieve_vocabulary_mismatch() {
2173 let tmp = tempdir().unwrap();
2174 let store = MemoryStore::open_default(tmp.path()).unwrap();
2175 let dog = sample_record(1, "My Golden Retriever loves the park");
2176 let other = sample_record(2, "I bought a yellow dress");
2177 let dog_id = dog.id;
2178 store.episodic().append_batch(&[dog, other]).unwrap();
2179 let hits = store
2180 .episodic()
2181 .search("What breed is my dog?", 5, false)
2182 .unwrap();
2183 assert_eq!(hits[0].record.id, dog_id);
2184 }
2185
2186 #[test]
2187 fn planner_boosts_temporal_date_match() {
2188 let tmp = tempdir().unwrap();
2189 let store = MemoryStore::open_default(tmp.path()).unwrap();
2190 let dated = sample_record(1, "I volunteered on February 14th at the animal shelter");
2191 let other = sample_record(2, "I volunteered at the community garden last summer");
2192 let dated_id = dated.id;
2193 store.episodic().append_batch(&[dated, other]).unwrap();
2194 let hits = store
2195 .episodic()
2196 .search("When did I volunteer at the animal shelter?", 5, false)
2197 .unwrap();
2198 assert_eq!(hits[0].record.id, dated_id);
2199 }
2200
2201 #[test]
2202 #[ignore = "manual in-process latency profile"]
2203 fn profile_ingest_and_search_latency() {
2204 fn timed_ms(label: &str, repeats: u32, mut work: impl FnMut()) {
2205 let start = std::time::Instant::now();
2206 for _ in 0..repeats {
2207 work();
2208 }
2209 let elapsed = start.elapsed();
2210 println!(
2211 "{label}: {:.3} ms (n={repeats})",
2212 elapsed.as_secs_f64() * 1000.0 / f64::from(repeats)
2213 );
2214 }
2215
2216 let search_records: Vec<EpisodicRecord> = (0..10_000)
2217 .map(|n| {
2218 sample_record(
2219 n,
2220 if n % 5 == 0 {
2221 "Rust memory retrieval benchmark item"
2222 } else {
2223 "Unrelated episodic record"
2224 },
2225 )
2226 })
2227 .collect();
2228
2229 timed_ms("append_single_1000", 1, || {
2230 let tmp = tempdir().unwrap();
2231 let store = MemoryStore::open_default(tmp.path()).unwrap();
2232 for n in 0..1_000 {
2233 store.episodic().append(&sample_record(n, "once")).unwrap();
2234 }
2235 });
2236 timed_ms("append_batch_1000", 1, || {
2237 let tmp = tempdir().unwrap();
2238 let store = MemoryStore::open_default(tmp.path()).unwrap();
2239 let records: Vec<EpisodicRecord> =
2240 (0..1_000).map(|n| sample_record(n, "once")).collect();
2241 store.episodic().append_batch(&records).unwrap();
2242 });
2243
2244 let tmp = tempdir().unwrap();
2245 {
2246 let store = MemoryStore::open_default(tmp.path()).unwrap();
2247 store.episodic().append_batch(&search_records).unwrap();
2248 }
2249 let cold = MemoryStore::open_default(tmp.path()).unwrap();
2250 timed_ms("cold_search_10000", 1, || {
2251 let hits = cold
2252 .episodic()
2253 .search("rust memory retrieval", 10, false)
2254 .unwrap();
2255 assert!(!hits.is_empty());
2256 });
2257 timed_ms("warm_search_10000", 50, || {
2258 let hits = cold
2259 .episodic()
2260 .search("rust memory retrieval", 10, false)
2261 .unwrap();
2262 assert!(!hits.is_empty());
2263 });
2264 }
2265
2266 #[test]
2272 #[ignore = "manual in-process latency profile at realistic scale"]
2273 fn profile_search_latency_25k_realistic() {
2274 const TOTAL: u64 = 25_000;
2275 const SESSIONS: u64 = 50;
2276 let topics = [
2277 "bookshelf",
2278 "guitar",
2279 "vegetarian",
2280 "portfolio",
2281 "commute",
2282 "grandmother",
2283 "chemistry",
2284 "marathon",
2285 "internship",
2286 "yoga",
2287 "spam filter",
2288 "projector",
2289 "swimming",
2290 "cousin",
2291 "bank account",
2292 "book club",
2293 "recipe",
2294 "journal subscription",
2295 "laptop",
2296 "hiking",
2297 ];
2298 let fillers = [
2299 "We discussed the plan for the weekend and agreed on the schedule.",
2300 "The meeting notes were circulated and everyone acknowledged them.",
2301 "I explained my reasoning and the group considered the proposal.",
2302 "After the presentation we reviewed the feedback together.",
2303 "She mentioned the deadline and we adjusted the timeline accordingly.",
2304 ];
2305
2306 let records: Vec<EpisodicRecord> = (0..TOTAL)
2307 .map(|n| {
2308 let session = n * SESSIONS / TOTAL;
2309 let topic = topics[(n as usize) % topics.len()];
2310 let filler = fillers[(n as usize) % fillers.len()];
2311 let content = format!(
2312 "Session {session} note {n}: my friend Alice mentioned {topic} while {filler}"
2313 );
2314 let session_id = if n % 4 == 0 {
2315 None
2316 } else {
2317 Some(uuid::Uuid::new_v4())
2318 };
2319 EpisodicRecord::new(
2320 session_id,
2321 n,
2322 if n % 3 == 0 {
2323 EpisodicKind::UserStatement
2324 } else {
2325 EpisodicKind::AssistantResponse
2326 },
2327 content,
2328 Provenance::new(ProvenanceSource::User),
2329 )
2330 })
2331 .collect();
2332
2333 let tmp = tempdir().unwrap();
2334 {
2335 let store = MemoryStore::open_default(tmp.path()).unwrap();
2336 let t = std::time::Instant::now();
2337 store.episodic().append_batch(&records).unwrap();
2338 println!("ingest 25k: {:.1} ms", t.elapsed().as_secs_f64() * 1000.0);
2339 }
2340
2341 let cold = MemoryStore::open_default(tmp.path()).unwrap();
2343 let episodic = cold.episodic();
2344 let t = std::time::Instant::now();
2345 let hits = episodic
2346 .search("guitar grandmother recipe", 10, false)
2347 .unwrap();
2348 println!(
2349 "cold_search: {:.1} ms (hits={})",
2350 t.elapsed().as_secs_f64() * 1000.0,
2351 hits.len()
2352 );
2353
2354 let queries: Vec<String> = (0..20)
2357 .map(|i| {
2358 let a = topics[i * 7 % topics.len()];
2359 let b = topics[(i * 7 + 5) % topics.len()];
2360 format!("{a} {b} weekend plan")
2361 })
2362 .collect();
2363 let start = std::time::Instant::now();
2364 for q in &queries {
2365 let hits = episodic.search(q, 10, false).unwrap();
2366 assert!(!hits.is_empty(), "no hits for {q}");
2367 }
2368 println!(
2369 "warm_search varied p50: {:.2} ms/query",
2370 start.elapsed().as_secs_f64() * 1000.0 / queries.len() as f64
2371 );
2372
2373 let t = std::time::Instant::now();
2374 let hits = episodic
2375 .search("zzzterm zzzother zzzthird", 10, false)
2376 .unwrap();
2377 println!(
2378 "no_match_query (full-scan fallback path): {:.1} ms (hits={})",
2379 t.elapsed().as_secs_f64() * 1000.0,
2380 hits.len()
2381 );
2382 }
2383
2384 #[test]
2385 fn enrichment_bridges_vocabulary_gap_for_theater() {
2386 let tmp = tempdir().unwrap();
2387 let store = MemoryStore::open_default(tmp.path()).unwrap();
2388 store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
2389 let answer = sample_record(1, "The production I attended was The Glass Menagerie");
2391 let competing = sample_record(2, "I went to a play at the local community theater");
2392 let answer_id = answer.id;
2393 store.episodic().append_batch(&[answer, competing]).unwrap();
2394 let hits = store
2395 .episodic()
2396 .search(
2397 "What play did I attend at the local community theater?",
2398 5,
2399 false,
2400 )
2401 .unwrap();
2402 assert!(hits.iter().any(|h| h.record.id == answer_id));
2405 }
2406
2407 #[test]
2408 fn enrichment_bridges_vocabulary_gap_for_shelter() {
2409 let tmp = tempdir().unwrap();
2410 let store = MemoryStore::open_default(tmp.path()).unwrap();
2411 store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
2412 let answer = sample_record(1, "I rescued a dog from the humane society last week");
2413 let other = sample_record(2, "I bought groceries at the store");
2414 let answer_id = answer.id;
2415 store.episodic().append_batch(&[answer, other]).unwrap();
2416 let hits = store
2417 .episodic()
2418 .search("When did I volunteer at the animal shelter?", 5, false)
2419 .unwrap();
2420 assert!(hits.iter().any(|h| h.record.id == answer_id));
2421 }
2422
2423 #[test]
2424 fn session_boost_favors_sessions_with_multiple_matches() {
2425 let tmp = tempdir().unwrap();
2426 let store = MemoryStore::open_default(tmp.path()).unwrap();
2427 let session_a = uuid::Uuid::new_v4();
2428 let session_b = uuid::Uuid::new_v4();
2429 let a1 = EpisodicRecord::new(
2431 Some(session_a),
2432 1,
2433 EpisodicKind::Observation,
2434 "I love hiking in the mountains",
2435 Provenance::new(ProvenanceSource::User),
2436 );
2437 let a2 = EpisodicRecord::new(
2438 Some(session_a),
2439 2,
2440 EpisodicKind::Observation,
2441 "Hiking in the mountains is great exercise",
2442 Provenance::new(ProvenanceSource::User),
2443 );
2444 let b1 = EpisodicRecord::new(
2446 Some(session_b),
2447 1,
2448 EpisodicKind::Observation,
2449 "Hiking is fun",
2450 Provenance::new(ProvenanceSource::User),
2451 );
2452 store.episodic().append_batch(&[a1, a2, b1]).unwrap();
2453 let hits = store
2454 .episodic()
2455 .search("hiking mountains", 10, false)
2456 .unwrap();
2457 let a_ranks: Vec<usize> = hits
2460 .iter()
2461 .enumerate()
2462 .filter(|(_, h)| h.record.session_id == Some(session_a))
2463 .map(|(i, _)| i)
2464 .collect();
2465 let b_rank = hits
2466 .iter()
2467 .position(|h| h.record.session_id == Some(session_b));
2468 if let Some(br) = b_rank {
2469 assert!(a_ranks.iter().all(|&ar| ar < br));
2470 }
2471 }
2472
2473 #[test]
2478 fn replace_content_rewrites_raw_record_and_term_postings() {
2479 let tmp = tempdir().unwrap();
2480 let store = MemoryStore::open_default(tmp.path()).unwrap();
2481 let secret = "zq8x5v3p1k9r2m4n6w7a2s4d8";
2484 let record = sample_record(1, &format!("TOKEN={secret}")).with_id(uuid::Uuid::new_v4());
2485 let id = record.id;
2486 store.episodic().append(&record).unwrap();
2487 assert!(
2488 !store
2489 .episodic()
2490 .search(secret, 5, false)
2491 .unwrap()
2492 .is_empty(),
2493 "fixture secret must be searchable before redaction"
2494 );
2495
2496 let redacted = "TOKEN=[REDACTED:credential_assignment]";
2497 assert!(store.episodic().replace_content(id, redacted).unwrap());
2498 let stored = store.episodic().get(id).unwrap().unwrap();
2499 assert_eq!(stored.content, redacted, "raw lane must be rewritten");
2500 assert!(
2501 store
2502 .episodic()
2503 .search(secret, 5, false)
2504 .unwrap()
2505 .is_empty(),
2506 "the old term must not keep the record searchable"
2507 );
2508 assert!(
2509 !store
2510 .episodic()
2511 .search("REDACTED", 5, false)
2512 .unwrap()
2513 .is_empty(),
2514 "the new terms must be posted"
2515 );
2516
2517 assert!(store.episodic().replace_content(id, redacted).unwrap());
2519 assert!(
2520 !store
2521 .episodic()
2522 .replace_content(uuid::Uuid::new_v4(), "x")
2523 .unwrap()
2524 );
2525 }
2526}