use lmdb::{Cursor, Database, Environment, Transaction, WriteFlags};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use wm_core::{
CoreError, EpisodicCapturePolicy, EpisodicId, EpisodicKind, EpisodicRecord, MemoryTransition,
Result, ValidityState,
};
use crate::embedder::Embedder;
use crate::enrichment::VocabularyEnrichment;
use crate::episodic_keys::{AdaptiveAliases, key_index_terms_with_aliases};
use crate::query_planner::QueryPlan;
use crate::search::strip_stopwords;
#[cfg(test)]
#[derive(Clone, Copy, PartialEq, Eq)]
enum CommitBoundary {
BeforeRawCommit,
AfterRawCommit,
}
#[cfg(test)]
fn commit_boundary_test_hook(boundary: CommitBoundary, records: &[EpisodicRecord]) -> Result<()> {
use std::io::{Read, Write};
let Ok(case) = std::env::var("WM_Q06_CASE") else {
return Ok(());
};
let expected = match boundary {
CommitBoundary::BeforeRawCommit => "before_raw_commit",
CommitBoundary::AfterRawCommit => "after_raw_commit",
};
if case != expected {
return Ok(());
}
let uuid = std::env::var("WM_Q06_UUID")
.map_err(|_| CoreError::Memory("q06 hook: WM_Q06_UUID is required".into()))?;
let names: Vec<String> = records.iter().map(|record| record.id.to_string()).collect();
if names.len() != 1 || names[0] != uuid {
return Err(CoreError::Memory(
"q06 hook: candidate set must be the single WM_Q06_UUID record".into(),
));
}
println!("WM_Q06_BOUNDARY {expected} {uuid}");
let _ = std::io::stdout().flush();
let mut release = [0_u8; 1];
match std::io::stdin().read(&mut release) {
Ok(0) => Err(CoreError::Memory(
"q06 hook: stdin closed before release".into(),
)),
Ok(_) => Ok(()),
Err(e) => Err(CoreError::Memory(format!(
"q06 hook: stdin read failed: {e}"
))),
}
}
#[derive(Debug, Clone)]
pub struct EpisodicSearchResult {
pub record: EpisodicRecord,
pub score: f32,
pub matched_terms: usize,
}
pub struct EpisodicStore<'a> {
env: &'a Environment,
db: Database,
term_db: Database,
term_cache: Arc<RwLock<HashMap<String, Vec<EpisodicId>>>>,
mutation_count: &'a std::sync::atomic::AtomicU64,
embedder: Option<Arc<dyn Embedder + Send + Sync>>,
aliases: Option<AdaptiveAliases>,
enrichment: Option<VocabularyEnrichment>,
}
impl<'a> EpisodicStore<'a> {
pub(crate) fn new(
env: &'a Environment,
db: Database,
term_db: Database,
term_cache: Arc<RwLock<HashMap<String, Vec<EpisodicId>>>>,
mutation_count: &'a std::sync::atomic::AtomicU64,
) -> Self {
Self {
env,
db,
term_db,
term_cache,
mutation_count,
embedder: None,
aliases: None,
enrichment: None,
}
}
#[must_use]
pub fn with_adaptive_aliases(mut self, aliases: AdaptiveAliases) -> Self {
if !aliases.is_empty() {
self.aliases = Some(aliases);
}
self
}
#[must_use]
pub fn with_enrichment(mut self, enrichment: VocabularyEnrichment) -> Self {
if !enrichment.is_empty() {
self.enrichment = Some(enrichment);
}
self
}
#[must_use]
pub fn with_embedder(mut self, embedder: Arc<dyn Embedder + Send + Sync>) -> Self {
self.embedder = Some(embedder);
self
}
pub fn append(&self, record: &EpisodicRecord) -> Result<()> {
self.append_batch(std::slice::from_ref(record))
}
pub fn append_batch(&self, records: &[EpisodicRecord]) -> Result<()> {
if records.is_empty() {
return Ok(());
}
let serialized = records
.iter()
.map(|record| {
rmp_serde::to_vec(record)
.map(|value| (record, value))
.map_err(|e| CoreError::Memory(format!("episodic serialize failed: {e}")))
})
.collect::<Result<Vec<_>>>()?;
let mut tx = self
.env
.begin_rw_txn()
.map_err(|e| CoreError::Memory(format!("episodic rw_txn failed: {e}")))?;
for (record, value) in &serialized {
match tx.put(
self.db,
record.id.as_bytes(),
value,
WriteFlags::NO_OVERWRITE,
) {
Ok(()) => {}
Err(lmdb::Error::KeyExist) => {
tx.abort();
return Err(CoreError::InvalidArgs(format!(
"episodic record {} already exists",
record.id
)));
}
Err(e) => {
tx.abort();
return Err(CoreError::Memory(format!("episodic append failed: {e}")));
}
}
}
#[cfg(test)]
commit_boundary_test_hook(CommitBoundary::BeforeRawCommit, records)?;
tx.commit()
.map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
#[cfg(test)]
commit_boundary_test_hook(CommitBoundary::AfterRawCommit, records)?;
self.mutation_count
.fetch_add(records.len() as u64, std::sync::atomic::Ordering::Relaxed);
self.index_records(records)?;
self.clear_term_cache();
Ok(())
}
fn term_postings(&self, term: &str) -> Result<Vec<EpisodicId>> {
let term = index_safe_term(term);
if let Ok(cache) = self.term_cache.read() {
if let Some(ids) = cache.get(&term) {
return Ok(ids.clone());
}
}
let tx = self
.env
.begin_ro_txn()
.map_err(|e| CoreError::Memory(format!("episodic index ro_txn failed: {e}")))?;
let term_key = term;
let mut ids: Vec<EpisodicId> = Vec::new();
{
match tx.get(self.term_db, &term_key) {
Ok(_) => {}
Err(lmdb::Error::NotFound) => {
tx.commit().map_err(|e| {
CoreError::Memory(format!("episodic index commit failed: {e}"))
})?;
if let Ok(mut cache) = self.term_cache.write() {
cache.insert(term_key, Vec::new());
}
return Ok(ids);
}
Err(e) => {
return Err(CoreError::Memory(format!(
"episodic index read failed: {e}"
)));
}
}
let mut cursor = tx
.open_ro_cursor(self.term_db)
.map_err(|e| CoreError::Memory(format!("episodic index cursor failed: {e}")))?;
for (key, value) in cursor.iter_from(term_key.as_bytes()) {
if key != term_key.as_bytes() {
break;
}
if value.len() == std::mem::size_of::<EpisodicId>() {
if let Ok(id) = EpisodicId::from_slice(value) {
ids.push(id);
}
}
}
}
tx.commit()
.map_err(|e| CoreError::Memory(format!("episodic index commit failed: {e}")))?;
if let Ok(mut cache) = self.term_cache.write() {
cache.insert(term_key, ids.clone());
}
Ok(ids)
}
fn clear_term_cache(&self) {
if let Ok(mut cache) = self.term_cache.write() {
cache.clear();
}
}
fn index_records(&self, records: &[EpisodicRecord]) -> Result<()> {
let public: Vec<&EpisodicRecord> = records
.iter()
.filter(|record| !record.is_private && !record.model_exclude)
.collect();
if public.is_empty() {
return Ok(());
}
let mut pending: HashMap<String, Vec<&EpisodicRecord>> = HashMap::new();
for record in &public {
let base_terms = index_terms_with_aliases(&record.content, self.aliases.as_ref());
let enriched: Vec<String> = if let Some(ref enrichment) = self.enrichment {
let mut all = base_terms.clone();
let extra = enrichment.enrich(&base_terms);
all.extend(extra);
all.sort();
all.dedup();
all
} else {
base_terms
};
for term in enriched {
pending
.entry(index_safe_term(&term))
.or_default()
.push(record);
}
}
let mut tx = self
.env
.begin_rw_txn()
.map_err(|e| CoreError::Memory(format!("episodic index rw_txn failed: {e}")))?;
for (term, records_for_term) in pending {
for record in records_for_term {
if let Err(e) = tx.put(
self.term_db,
&term,
&record.id.as_bytes(),
WriteFlags::default(),
) {
tx.abort();
return Err(CoreError::Memory(format!(
"episodic term index write failed: {e}"
)));
}
}
}
tx.commit()
.map_err(|e| CoreError::Memory(format!("episodic term index commit failed: {e}")))?;
Ok(())
}
pub fn rebuild_sidecar(&self) -> Result<usize> {
let records = self.scan(None, usize::MAX)?;
let mut indexed = 0usize;
for chunk in records.chunks(5_000) {
self.index_records(chunk)?;
indexed += chunk.len();
}
self.clear_term_cache();
Ok(indexed)
}
pub fn sidecar_is_empty(&self) -> Result<bool> {
let tx = self
.env
.begin_ro_txn()
.map_err(|e| CoreError::Memory(format!("episodic index ro_txn failed: {e}")))?;
let mut cursor = tx
.open_ro_cursor(self.term_db)
.map_err(|e| CoreError::Memory(format!("episodic index cursor failed: {e}")))?;
Ok(cursor.iter().next().is_none())
}
pub fn record_count(&self) -> Result<u64> {
let tx = self
.env
.begin_ro_txn()
.map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
let mut cursor = tx
.open_ro_cursor(self.db)
.map_err(|e| CoreError::Memory(format!("episodic cursor failed: {e}")))?;
let mut count = 0u64;
for _ in cursor.iter() {
count += 1;
}
Ok(count)
}
pub fn append_explicit(
&self,
record: &EpisodicRecord,
policy: EpisodicCapturePolicy,
) -> Result<bool> {
let prepared = record
.clone()
.with_content(policy.prepare_content(&record.content));
self.append(&prepared)?;
Ok(true)
}
pub fn append_explicit_batch(
&self,
records: &[EpisodicRecord],
policy: EpisodicCapturePolicy,
) -> Result<usize> {
if records.is_empty() {
return Ok(0);
}
let prepared: Vec<EpisodicRecord> = records
.iter()
.map(|record| {
record
.clone()
.with_content(policy.prepare_content(&record.content))
})
.collect();
self.append_batch(&prepared)?;
Ok(prepared.len())
}
pub fn get(&self, id: EpisodicId) -> Result<Option<EpisodicRecord>> {
let tx = self
.env
.begin_ro_txn()
.map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
let result = tx.get(self.db, id.as_bytes());
match result {
Ok(bytes) => {
let record: EpisodicRecord = rmp_serde::from_slice(bytes)
.map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
tx.commit()
.map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
Ok(Some(record))
}
Err(lmdb::Error::NotFound) => {
tx.commit()
.map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
Ok(None)
}
Err(e) => Err(CoreError::Memory(format!("episodic get failed: {e}"))),
}
}
pub fn transition(&self, id: EpisodicId, transition: MemoryTransition) -> Result<()> {
let mut tx = self
.env
.begin_rw_txn()
.map_err(|e| CoreError::Memory(format!("episodic rw_txn failed: {e}")))?;
let bytes = match tx.get(self.db, id.as_bytes()) {
Ok(bytes) => bytes,
Err(lmdb::Error::NotFound) => {
tx.abort();
return Err(CoreError::InvalidArgs(format!(
"episodic record {id} does not exist"
)));
}
Err(e) => {
tx.abort();
return Err(CoreError::Memory(format!("episodic get failed: {e}")));
}
};
let mut record: EpisodicRecord = rmp_serde::from_slice(bytes)
.map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
record
.transition(transition)
.map_err(|e| CoreError::InvalidArgs(format!("episodic transition rejected: {e}")))?;
let value = rmp_serde::to_vec(&record)
.map_err(|e| CoreError::Memory(format!("episodic serialize failed: {e}")))?;
tx.put(self.db, id.as_bytes(), &value, WriteFlags::default())
.map_err(|e| CoreError::Memory(format!("episodic transition write failed: {e}")))?;
tx.commit()
.map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
self.mutation_count
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
Ok(())
}
pub fn scan(
&self,
session_id: Option<uuid::Uuid>,
limit: usize,
) -> Result<Vec<EpisodicRecord>> {
if limit == 0 {
return Ok(Vec::new());
}
let tx = self
.env
.begin_ro_txn()
.map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
let mut cursor = tx
.open_ro_cursor(self.db)
.map_err(|e| CoreError::Memory(format!("episodic cursor failed: {e}")))?;
let mut records = Vec::new();
for item in cursor.iter() {
let (_, bytes) = item;
let record: EpisodicRecord = rmp_serde::from_slice(bytes)
.map_err(|e| CoreError::Memory(format!("episodic deserialize failed: {e}")))?;
if session_id.is_none_or(|id| record.session_id == Some(id)) {
records.push(record);
}
}
drop(cursor);
tx.commit()
.map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
records.sort_by_key(|record| (record.sequence, record.created_at, record.id));
records.truncate(limit);
Ok(records)
}
pub fn search(
&self,
query: &str,
limit: usize,
include_historical: bool,
) -> Result<Vec<EpisodicSearchResult>> {
self.search_with_limits(query, limit, limit.saturating_mul(2), include_historical)
}
pub fn search_with_limits(
&self,
query: &str,
limit: usize,
candidate_limit: usize,
include_historical: bool,
) -> Result<Vec<EpisodicSearchResult>> {
if limit == 0 {
return Ok(Vec::new());
}
let mut results = self.search_scored(query, limit, candidate_limit, include_historical)?;
if is_current_query(query) {
self.resolve_current(&mut results);
}
results.truncate(limit);
Ok(results)
}
fn search_scored(
&self,
query: &str,
limit: usize,
candidate_limit: usize,
include_historical: bool,
) -> Result<Vec<EpisodicSearchResult>> {
let plan = QueryPlan::plan(query, limit);
let candidate_limit = candidate_limit.max(plan.candidate_limit);
let query_terms = tokenize(query);
let query_keys = key_index_terms_with_aliases(query, self.aliases.as_ref());
if query_terms.is_empty() && query_keys.is_empty() || limit == 0 {
return Ok(Vec::new());
}
let mut candidate_scores: HashMap<EpisodicId, usize> = HashMap::new();
for term in query_terms.iter().chain(query_keys.iter()) {
for id in self.term_postings(term)? {
*candidate_scores.entry(id).or_default() += 1;
}
}
let records = if candidate_scores.is_empty() {
if !matches!(self.sidecar_is_empty(), Ok(true)) {
return Ok(Vec::new());
}
self.scan(None, usize::MAX)?
} else {
let mut ranked_candidates: Vec<(EpisodicId, usize)> =
candidate_scores.into_iter().collect();
ranked_candidates.sort_by(|(left_id, left_count), (right_id, right_count)| {
right_count
.cmp(left_count)
.then_with(|| left_id.cmp(right_id))
});
ranked_candidates.truncate(candidate_limit);
self.load_records(
&ranked_candidates
.into_iter()
.map(|(id, _)| id)
.collect::<Vec<_>>(),
)?
};
let mut results = Vec::new();
for record in records {
if !include_historical && !matches!(record.validity, ValidityState::Active) {
continue;
}
let content_terms = tokenize(&record.content);
let content_keys = key_index_terms_with_aliases(&record.content, self.aliases.as_ref());
let reverse_map: HashMap<&String, Vec<String>> =
if let Some(ref enrichment) = self.enrichment {
if matches!(record.kind, EpisodicKind::UserStatement) {
query_terms
.iter()
.map(|qt| (qt, enrichment.reverse_enrich(qt)))
.collect()
} else {
HashMap::new()
}
} else {
HashMap::new()
};
let mut reverse_match_count = 0usize;
let matched_terms = query_terms
.iter()
.filter(|term| {
if content_terms.iter().any(|candidate| candidate == *term)
|| content_keys.iter().any(|candidate| candidate == *term)
{
return true;
}
if let Some(reverse_terms) = reverse_map.get(term) {
let found = reverse_terms.iter().any(|rt| {
content_terms.iter().any(|candidate| candidate == rt)
|| content_keys.iter().any(|candidate| candidate == rt)
});
if found {
reverse_match_count += 1;
}
return found;
}
false
})
.count();
let matched_keys = query_keys
.iter()
.filter(|term| {
content_keys.iter().any(|candidate| candidate == *term)
|| content_terms.iter().any(|candidate| candidate == *term)
})
.count();
if matched_terms == 0 && matched_keys == 0 {
continue;
}
let key_bonus = if query_keys.is_empty() {
0.0
} else {
matched_keys as f32 / query_keys.len() as f32 * plan.key_weight
};
let role_boost = match record.kind {
EpisodicKind::UserStatement => 0.12,
_ => 0.0,
};
let effective_matched = if matches!(record.kind, EpisodicKind::UserStatement) {
(matched_terms + 2).min(query_terms.len())
} else {
matched_terms
};
let coverage = if query_terms.is_empty() {
0.0
} else {
effective_matched as f32 / query_terms.len() as f32
};
let number_bonus = if plan.number_query {
let has_digit = content_terms
.iter()
.any(|term| term.chars().any(|c| c.is_ascii_digit()));
if has_digit || contains_number_word(&record.content) {
0.03
} else {
0.0
}
} else {
0.0
};
let density = matched_terms as f32 / content_terms.len().max(1) as f32;
results.push(EpisodicSearchResult {
record,
score: coverage
+ key_bonus
+ role_boost
+ number_bonus
+ (reverse_match_count as f32).mul_add(0.05, density * 0.03),
matched_terms: matched_terms.max(matched_keys),
});
}
let mut session_counts: HashMap<Option<uuid::Uuid>, usize> = HashMap::new();
for r in &results {
*session_counts.entry(r.record.session_id).or_default() += 1;
}
for r in &mut results {
let count = session_counts
.get(&r.record.session_id)
.copied()
.unwrap_or(1);
if count > 1 {
r.score = 0.02f32.mul_add((count - 1).min(3) as f32, r.score);
}
}
let mut hash_counts: HashMap<&str, usize> = HashMap::new();
for r in &results {
*hash_counts
.entry(r.record.content_hash.as_str())
.or_default() += 1;
}
let hash_boosts: HashMap<String, f32> = results
.iter()
.map(|r| {
let count = hash_counts
.get(r.record.content_hash.as_str())
.copied()
.unwrap_or(1);
let boost = if count > 1 {
0.03 * (count - 1).min(3) as f32
} else {
0.0
};
(r.record.id.to_string(), boost)
})
.collect();
for r in &mut results {
if let Some(boost) = hash_boosts.get(&r.record.id.to_string()) {
r.score += boost;
}
}
results.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| b.matched_terms.cmp(&a.matched_terms))
.then_with(|| a.record.content.len().cmp(&b.record.content.len()))
.then_with(|| a.record.sequence.cmp(&b.record.sequence))
.then_with(|| a.record.id.cmp(&b.record.id))
});
Ok(results)
}
pub fn search_with_rerank(
&self,
query: &str,
limit: usize,
candidate_limit: usize,
include_historical: bool,
alpha: f32,
rerank_pool: usize,
) -> Result<Vec<EpisodicSearchResult>> {
let Some(ref embedder) = self.embedder else {
return self.search_with_limits(query, limit, candidate_limit, include_historical);
};
if !embedder.is_available() || limit == 0 {
return self.search_with_limits(query, limit, candidate_limit, include_historical);
}
let rerank_pool = if rerank_pool == 0 {
limit.max(candidate_limit).min(50)
} else {
rerank_pool.max(limit).min(50)
};
let deterministic: Vec<EpisodicSearchResult> = self
.search_scored(query, rerank_pool, candidate_limit, include_historical)?
.into_iter()
.take(rerank_pool)
.collect();
if deterministic.is_empty() {
return Ok(Vec::new());
}
let contents: Vec<&str> = std::iter::once(query)
.chain(deterministic.iter().map(|r| r.record.content.as_str()))
.collect();
let embeddings = embedder.embed_batch(&contents)?;
if embeddings.len() != deterministic.len() + 1 {
return Err(CoreError::Memory(format!(
"embedder returned {} vectors, expected {}",
embeddings.len(),
deterministic.len() + 1
)));
}
let query_vec = &embeddings[0];
let candidate_vecs = &embeddings[1..];
if alpha >= 2.0 {
let protected: Vec<EpisodicSearchResult> =
deterministic.into_iter().take(limit).collect();
let cosines: Vec<f32> = protected
.iter()
.enumerate()
.map(|(i, _)| cosine_sim(query_vec, &candidate_vecs[i]))
.collect();
let mut order: Vec<usize> = (0..protected.len()).collect();
order.sort_by(|&a, &b| {
cosines[b]
.partial_cmp(&cosines[a])
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.cmp(&b))
});
let mut slots: Vec<Option<EpisodicSearchResult>> =
protected.into_iter().map(Some).collect();
let reranked: Vec<EpisodicSearchResult> =
order.into_iter().filter_map(|i| slots[i].take()).collect();
Ok(reranked)
} else if alpha >= 1.0 {
let delta = 0.05;
let mut reranked = deterministic;
let cosines: Vec<f32> = candidate_vecs
.iter()
.map(|v| cosine_sim(query_vec, v))
.collect();
let n = reranked.len();
for _ in 0..n {
let mut swapped = false;
for i in 0..n.saturating_sub(1) {
let det_gap = (reranked[i].score - reranked[i + 1].score).abs();
if det_gap < delta && cosines[i + 1] > cosines[i] {
reranked.swap(i, i + 1);
swapped = true;
}
}
if !swapped {
break;
}
}
if is_current_query(query) {
self.resolve_current(&mut reranked);
}
reranked.truncate(limit);
Ok(reranked)
} else {
let max_det = deterministic
.iter()
.map(|r| r.score)
.fold(0.0f32, f32::max)
.max(1e-9);
let mut reranked: Vec<EpisodicSearchResult> = deterministic
.into_iter()
.enumerate()
.map(|(i, mut r)| {
let cosine = cosine_sim(query_vec, &candidate_vecs[i]);
let det_norm = r.score / max_det;
r.score = alpha.mul_add(det_norm, (1.0 - alpha) * cosine);
r
})
.collect();
reranked.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| b.matched_terms.cmp(&a.matched_terms))
.then_with(|| a.record.content.len().cmp(&b.record.content.len()))
.then_with(|| a.record.sequence.cmp(&b.record.sequence))
.then_with(|| a.record.id.cmp(&b.record.id))
});
if is_current_query(query) {
self.resolve_current(&mut reranked);
}
reranked.truncate(limit);
Ok(reranked)
}
}
fn resolve_current(&self, results: &mut Vec<EpisodicSearchResult>) {
if results.len() < 2 {
return;
}
let mut anchors: Vec<EpisodicSearchResult> = Vec::new();
let mut rest: Vec<EpisodicSearchResult> = Vec::new();
for result in results.drain(..) {
let is_anchor = matches!(result.record.kind, EpisodicKind::UserStatement)
&& contains_change_marker(&result.record.content);
if is_anchor {
anchors.push(result);
} else {
rest.push(result);
}
}
if anchors.is_empty() {
*results = rest;
return;
}
anchors.sort_by(|a, b| {
b.record
.created_at
.cmp(&a.record.created_at)
.then_with(|| b.record.sequence.cmp(&a.record.sequence))
.then_with(|| a.record.id.cmp(&b.record.id))
});
anchors.extend(rest);
*results = anchors;
}
fn load_records(&self, ids: &[EpisodicId]) -> Result<Vec<EpisodicRecord>> {
let tx = self
.env
.begin_ro_txn()
.map_err(|e| CoreError::Memory(format!("episodic ro_txn failed: {e}")))?;
let mut records = Vec::with_capacity(ids.len());
for id in ids {
match tx.get(self.db, id.as_bytes()) {
Ok(bytes) => {
records.push(rmp_serde::from_slice(bytes).map_err(|e| {
CoreError::Memory(format!("episodic deserialize failed: {e}"))
})?);
}
Err(lmdb::Error::NotFound) => {}
Err(e) => return Err(CoreError::Memory(format!("episodic get failed: {e}"))),
}
}
tx.commit()
.map_err(|e| CoreError::Memory(format!("episodic commit failed: {e}")))?;
Ok(records)
}
}
fn index_terms_with_aliases(text: &str, aliases: Option<&AdaptiveAliases>) -> Vec<String> {
tokenize(text)
.into_iter()
.chain(key_index_terms_with_aliases(text, aliases))
.fold(Vec::new(), |mut terms, term| {
if !terms.contains(&term) {
terms.push(term);
}
terms
})
}
const MAX_TERM_KEY_BYTES: usize = 480;
fn index_safe_term(term: &str) -> String {
if term.len() <= MAX_TERM_KEY_BYTES {
return term.to_string();
}
let digest = Sha256::digest(term.as_bytes());
format!("~h:{digest:x}")
}
fn tokenize(text: &str) -> Vec<String> {
strip_stopwords(text)
.split(|c: char| !c.is_alphanumeric())
.filter(|term| term.len() > 1)
.map(|term| simple_stem(&term.to_ascii_lowercase()))
.fold(Vec::new(), |mut terms, term| {
if !terms.contains(&term) {
terms.push(term);
}
terms
})
}
const CURRENT_QUERY_WORD_CUES: &[&str] = &["current", "currently", "latest", "nowadays"];
const CURRENT_QUERY_PHRASE_CUES: &[&str] = &["these days", "right now", "at the moment"];
#[must_use]
pub fn is_current_query(query: &str) -> bool {
let lowered = query.to_ascii_lowercase();
let has_word = lowered
.split(|c: char| !c.is_alphanumeric())
.any(|token| CURRENT_QUERY_WORD_CUES.contains(&token));
has_word
|| CURRENT_QUERY_PHRASE_CUES
.iter()
.any(|cue| lowered.contains(cue))
}
const CHANGE_MARKERS: &[&str] = &[
"switched to",
"switch to",
"switching to",
"switched from",
"changed my",
"change my",
"changed from",
"now prefer",
"now i prefer",
"now i'm",
"now im",
"no longer",
"used to",
"moved to",
"not anymore",
"instead of",
"replaced",
"gave up",
];
fn contains_change_marker(content: &str) -> bool {
let lowered = content.to_ascii_lowercase();
CHANGE_MARKERS.iter().any(|marker| lowered.contains(marker))
}
const CONTRADICTION_MARKERS: &[&str] = &[
"no longer",
"anymore",
"changed my mind",
"changed my",
"used to",
"gave up",
"just a phase",
"not really",
"but i",
];
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct EpisodicConflict {
pub later_record: EpisodicId,
pub earlier_record: EpisodicId,
pub marker: String,
pub shared_terms: Vec<String>,
pub later_content: String,
pub earlier_content: String,
}
#[must_use]
pub fn detect_conflicts(results: &[EpisodicSearchResult]) -> Vec<EpisodicConflict> {
const MAX_CONFLICTS: usize = 10;
let mut conflicts: Vec<EpisodicConflict> = Vec::new();
let mut seen_pairs: Vec<(EpisodicId, EpisodicId)> = Vec::new();
for (i, marked) in results.iter().enumerate() {
if !matches!(marked.record.kind, EpisodicKind::UserStatement) {
continue;
}
let lowered = marked.record.content.to_ascii_lowercase();
let Some(marker) = CONTRADICTION_MARKERS
.iter()
.find(|m| lowered.contains(*m))
.copied()
else {
continue;
};
let marked_terms = tokenize(&marked.record.content);
for (j, other) in results.iter().enumerate() {
if i == j || !matches!(other.record.kind, EpisodicKind::UserStatement) {
continue;
}
if other.record.content_hash == marked.record.content_hash {
continue;
}
let other_terms = tokenize(&other.record.content);
let shared: Vec<String> = marked_terms
.iter()
.filter(|t| other_terms.contains(t))
.cloned()
.collect();
if shared.len() < 2 {
continue;
}
let (later, earlier) = if (marked.record.created_at, marked.record.sequence)
> (other.record.created_at, other.record.sequence)
{
(&marked.record, &other.record)
} else {
(&other.record, &marked.record)
};
let pair_key = if later.id < earlier.id {
(later.id, earlier.id)
} else {
(earlier.id, later.id)
};
if seen_pairs.contains(&pair_key) {
continue;
}
seen_pairs.push(pair_key);
conflicts.push(EpisodicConflict {
later_record: later.id,
earlier_record: earlier.id,
marker: marker.to_string(),
shared_terms: shared,
later_content: later.content.clone(),
earlier_content: earlier.content.clone(),
});
if conflicts.len() >= MAX_CONFLICTS {
return conflicts;
}
}
}
conflicts
}
fn simple_stem(word: &str) -> String {
if word.len() <= 3 {
return word.to_string();
}
for suffix in ["ies", "ied", "ing", "edly", "ed", "ly", "es", "s"] {
if let Some(stem) = word.strip_suffix(suffix) {
if suffix == "ies" || suffix == "ied" {
return format!("{stem}y");
}
if stem.len() >= 2 {
return stem.to_string();
}
}
}
word.to_string()
}
fn cosine_sim(a: &[f32], b: &[f32]) -> f32 {
let dot = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum::<f32>();
let norm_a = a.iter().map(|x| x * x).sum::<f32>().sqrt();
let norm_b = b.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm_a < 1e-9 || norm_b < 1e-9 {
0.0
} else {
dot / (norm_a * norm_b)
}
}
fn contains_number_word(text: &str) -> bool {
const NUMBER_WORDS: &[&str] = &[
"one",
"two",
"three",
"four",
"five",
"six",
"seven",
"eight",
"nine",
"ten",
"eleven",
"twelve",
"thirteen",
"fourteen",
"fifteen",
"sixteen",
"seventeen",
"eighteen",
"nineteen",
"twenty",
"thirty",
"forty",
"fifty",
"sixty",
"seventy",
"eighty",
"ninety",
"hundred",
"thousand",
"million",
"billion",
"dozen",
"couple",
"half",
"quarter",
"double",
"triple",
"twice",
];
for word in text.split(|c: char| !c.is_alphanumeric()) {
if word.len() >= 3 && NUMBER_WORDS.iter().any(|nw| word.eq_ignore_ascii_case(nw)) {
return true;
}
}
false
}
#[cfg(test)]
mod tests {
use super::*;
use crate::MemoryStore;
#[cfg(unix)]
use chrono::{DateTime, Utc};
use tempfile::tempdir;
#[cfg(unix)]
use uuid::Uuid;
use wm_core::{EpisodicKind, Provenance, ProvenanceSource, ValidityState};
fn sample_record(sequence: u64, content: &str) -> EpisodicRecord {
EpisodicRecord::new(
None,
sequence,
EpisodicKind::Observation,
content,
Provenance::new(ProvenanceSource::User),
)
}
fn user_statement(sequence: u64, content: &str) -> EpisodicRecord {
EpisodicRecord::new(
None,
sequence,
EpisodicKind::UserStatement,
content,
Provenance::new(ProvenanceSource::User),
)
}
fn assistant_response(sequence: u64, content: &str) -> EpisodicRecord {
EpisodicRecord::new(
None,
sequence,
EpisodicKind::AssistantResponse,
content,
Provenance::new(ProvenanceSource::Agent),
)
}
#[cfg(unix)]
fn q06_record(id: u128, sequence: u64, content: &str, created_at: &str) -> EpisodicRecord {
let mut record = EpisodicRecord::new(
None,
sequence,
EpisodicKind::Observation,
content,
Provenance::new(ProvenanceSource::User),
)
.with_id(Uuid::from_u128(id));
record.created_at = DateTime::parse_from_rfc3339(created_at)
.unwrap()
.with_timezone(&Utc);
record
}
#[cfg(unix)]
fn q06_acknowledged() -> EpisodicRecord {
q06_record(601, 601, "q06 acknowledged control", "2026-01-01T00:10:01Z")
}
#[cfg(unix)]
fn run_q06_child(case: &str, store_path: &std::path::Path) {
use std::io::Write;
let expected_uuid = std::env::var("WM_Q06_UUID").expect("WM_Q06_UUID");
let uuid = Uuid::parse_str(&expected_uuid).expect("WM_Q06_UUID must parse");
let (sequence, content, created_at) = match case {
"before_raw_commit" => (602, "q06 precommit candidate", "2026-01-01T00:10:02Z"),
"after_raw_commit" => (603, "q06 uncertain candidate", "2026-01-01T00:10:03Z"),
other => panic!("q06 child: unknown case {other}"),
};
let record = q06_record(uuid.as_u128(), sequence, content, created_at);
let store = MemoryStore::open_default(store_path).expect("q06 child store");
match store.episodic().append(&record) {
Ok(()) => {
println!("WM_Q06_CALLER_ACK {uuid}");
let _ = std::io::stdout().flush();
}
Err(e) => {
eprintln!("q06 child append failed: {e}");
std::process::exit(2);
}
}
}
#[cfg(unix)]
fn run_q06_killed_case(
test_filter: &str,
store_path: &std::path::Path,
case: &str,
uuid: Uuid,
) -> Vec<String> {
use std::io::{BufRead, BufReader};
let exe = std::env::current_exe().expect("q06 current_exe");
let mut child = std::process::Command::new(exe)
.arg(test_filter)
.arg("--exact")
.arg("--nocapture")
.env("WM_Q06_CASE", case)
.env("WM_Q06_STORE", store_path)
.env("WM_Q06_UUID", uuid.to_string())
.stdout(std::process::Stdio::piped())
.stdin(std::process::Stdio::piped())
.stderr(std::process::Stdio::inherit())
.spawn()
.expect("q06 spawn child");
let stdout = child.stdout.take().expect("q06 child stdout");
let (tx, rx) = std::sync::mpsc::channel::<String>();
let reader = std::thread::spawn(move || {
for line in BufReader::new(stdout).lines() {
match line {
Ok(line) => {
if tx.send(line).is_err() {
break;
}
}
Err(_) => break,
}
}
});
let expected = format!("WM_Q06_BOUNDARY {case} {uuid}");
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
let mut lines = Vec::new();
let mut seen = false;
while !seen {
let remaining = deadline.saturating_duration_since(std::time::Instant::now());
if remaining.is_zero() {
break;
}
match rx.recv_timeout(remaining) {
Ok(line) if line == expected => seen = true,
Ok(line) => lines.push(line),
Err(_) => break,
}
}
if !seen {
let _ = child.kill();
let _ = child.wait();
reader.join().ok();
panic!("q06 case {case}: boundary {expected:?} not observed; lines={lines:?}");
}
child.kill().expect("q06 kill blocked child");
let status = child.wait().expect("q06 reap child");
assert!(
!status.success(),
"q06 case {case}: killed child must not exit successfully: {status:?}"
);
reader.join().ok();
while let Ok(line) = rx.try_recv() {
lines.push(line);
}
assert!(
!lines.iter().any(|line| line.contains("WM_Q06_CALLER_ACK")),
"q06 case {case}: caller acknowledgement in a killed case invalidates the experiment: {lines:?}"
);
lines
}
#[cfg(unix)]
#[test]
fn q06_commit_boundary_sigkill_classification() {
if let Ok(case) = std::env::var("WM_Q06_CASE") {
let store_path = std::env::var("WM_Q06_STORE").expect("WM_Q06_STORE");
run_q06_child(&case, std::path::Path::new(&store_path));
return;
}
let dir = tempdir().unwrap();
let store_path = dir.path().join("lmdb");
let acknowledged = q06_acknowledged();
{
let store = MemoryStore::open_default(&store_path).unwrap();
store.episodic().append(&acknowledged).unwrap();
let read = store
.episodic()
.get(acknowledged.id)
.unwrap()
.expect("acknowledged record");
assert_eq!(read, acknowledged);
}
let test_filter = "episodic::tests::q06_commit_boundary_sigkill_classification";
let candidate_before =
q06_record(602, 602, "q06 precommit candidate", "2026-01-01T00:10:02Z");
run_q06_killed_case(
test_filter,
&store_path,
"before_raw_commit",
candidate_before.id,
);
{
let store = MemoryStore::open_default(&store_path).unwrap();
assert_eq!(
store.episodic().get(acknowledged.id).unwrap(),
Some(acknowledged.clone()),
"acknowledged record must survive a pre-commit kill unchanged"
);
assert_eq!(
store.episodic().get(candidate_before.id).unwrap(),
None,
"pre-commit candidate must be absent after reopen"
);
}
let candidate_after =
q06_record(603, 603, "q06 uncertain candidate", "2026-01-01T00:10:03Z");
run_q06_killed_case(
test_filter,
&store_path,
"after_raw_commit",
candidate_after.id,
);
{
let store = MemoryStore::open_default(&store_path).unwrap();
assert_eq!(
store.episodic().get(acknowledged.id).unwrap(),
Some(acknowledged),
"acknowledged record must survive a post-commit kill unchanged"
);
assert_eq!(
store.episodic().get(candidate_before.id).unwrap(),
None,
"pre-commit candidate must stay absent"
);
assert_eq!(
store.episodic().get(candidate_after.id).unwrap(),
Some(candidate_after),
"post-commit candidate must be present and byte-equal after reopen"
);
}
}
#[test]
fn current_query_detection() {
assert!(is_current_query("What's my current favorite coffee?"));
assert!(is_current_query("What am I currently reading these days?"));
assert!(is_current_query("What's the latest book I mentioned?"));
assert!(is_current_query("What's my job right now?"));
assert!(is_current_query("What am I eating at the moment?"));
assert!(!is_current_query("What's my favorite coffee?"));
assert!(!is_current_query("Where did I volunteer in February?"));
assert!(!is_current_query("What did I say about the trip?"));
assert!(!is_current_query("What currency did I use in Japan?"));
}
#[test]
fn protected_rerank_preserves_membership_with_a_small_pool() {
use crate::embedder::Embedder;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
struct CountingEmbedder(Arc<AtomicUsize>);
impl Embedder for CountingEmbedder {
fn embed_batch(&self, texts: &[&str]) -> wm_core::Result<Vec<Vec<f32>>> {
self.0.store(texts.len(), Ordering::Relaxed);
Ok(texts
.iter()
.map(|t| {
let mut v = vec![0.0_f32; 8];
for (i, b) in t.bytes().enumerate() {
v[i % 8] += f32::from(b) / 255.0;
}
v
})
.collect())
}
fn dimension(&self) -> usize {
8
}
fn is_available(&self) -> bool {
true
}
fn backend_name(&self) -> &'static str {
"counting-test"
}
}
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
for i in 0..20 {
store
.episodic()
.append(&user_statement(
i,
&format!("Record {i} discusses topic alpha beta gamma delta epsilon"),
))
.unwrap();
}
let query = "topic alpha beta gamma delta epsilon";
let expected: Vec<uuid::Uuid> = store
.episodic()
.search_scored(query, 5, 20, false)
.unwrap()
.iter()
.take(5)
.map(|r| r.record.id)
.collect();
assert_eq!(expected.len(), 5);
let batches = Arc::new(AtomicUsize::new(0));
store.set_episodic_embedder(Arc::new(CountingEmbedder(batches.clone())));
let reranked = store
.episodic()
.search_with_rerank(query, 5, 20, false, 2.0, 5)
.unwrap();
assert_eq!(
batches.load(Ordering::Relaxed),
6,
"small rerank pool must embed pool + query only"
);
let mut got: Vec<uuid::Uuid> = reranked.iter().map(|r| r.record.id).collect();
let mut want = expected;
got.sort();
want.sort();
assert_eq!(
got, want,
"protected mode must keep the deterministic candidate set (membership), only reorder it"
);
}
#[test]
fn current_query_resolution_prefers_latest_statement() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = store.episodic();
episodic
.append(&user_statement(1, "My favorite coffee is dark roast."))
.unwrap();
episodic
.append(&user_statement(
2,
"I really love dark roast when it comes to coffee.",
))
.unwrap();
episodic
.append(&user_statement(3, "I've been jogging lately."))
.unwrap();
episodic
.append(&user_statement(4, "I've switched to cold brew for coffee."))
.unwrap();
let results = episodic
.search("What's my current favorite coffee?", 5, false)
.unwrap();
assert!(!results.is_empty());
assert!(
results[0].record.content.contains("cold brew"),
"current query must rank the latest statement first, got: {}",
results[0].record.content
);
}
#[test]
fn current_query_anchors_switched_from_template() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = store.episodic();
episodic
.append(&user_statement(1, "My favorite coffee is espresso."))
.unwrap();
episodic
.append(&user_statement(
2,
"I've actually switched from espresso to latte for coffee.",
))
.unwrap();
episodic
.append(&user_statement(3, "My favorite coffee is latte."))
.unwrap();
let results = episodic
.search("What's my current favorite coffee?", 5, false)
.unwrap();
assert!(!results.is_empty());
assert!(
results[0].record.content.contains("latte"),
"'switched from' must anchor the current value, got: {}",
results[0].record.content
);
}
#[test]
fn non_current_query_keeps_score_order() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = store.episodic();
episodic
.append(&user_statement(1, "My favorite coffee is dark roast."))
.unwrap();
episodic
.append(&user_statement(2, "I've switched to cold brew for coffee."))
.unwrap();
let results = episodic
.search("What's my favorite coffee?", 5, false)
.unwrap();
assert!(!results.is_empty());
assert!(
results[0].record.content.contains("dark roast"),
"non-current query must keep score order, got: {}",
results[0].record.content
);
}
#[test]
fn current_resolution_anchors_on_user_statements_only() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = store.episodic();
episodic
.append(&user_statement(1, "My favorite coffee is dark roast."))
.unwrap();
episodic
.append(&assistant_response(
2,
"Got it, dark roast is your favorite coffee!",
))
.unwrap();
episodic
.append(&user_statement(3, "I've switched to cold brew for coffee."))
.unwrap();
let results = episodic
.search("What's my current favorite coffee?", 5, false)
.unwrap();
assert!(
results[0].record.content.contains("cold brew"),
"user statements anchor chronology, got: {}",
results[0].record.content
);
assert_eq!(results[0].record.kind, EpisodicKind::UserStatement);
}
#[test]
fn current_query_without_change_markers_keeps_score_order() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = store.episodic();
episodic
.append(&user_statement(
1,
"My favorite hiking trail is Eagle Ridge.",
))
.unwrap();
episodic
.append(&user_statement(2, "I go hiking every weekend."))
.unwrap();
let results = episodic
.search("What's my current favorite hiking trail?", 5, false)
.unwrap();
assert!(!results.is_empty());
assert!(
results[0].record.content.contains("Eagle Ridge"),
"no change markers → deterministic score order, got: {}",
results[0].record.content
);
}
#[test]
fn detect_conflicts_flags_contradiction_with_shared_topic() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = store.episodic();
episodic
.append(&user_statement(
1,
"I'm vegetarian now. I decided to stop eating animal products.",
))
.unwrap();
episodic
.append(&user_statement(
2,
"I'm not really vegetarian anymore, I eat steak now.",
))
.unwrap();
episodic
.append(&user_statement(3, "I went hiking yesterday."))
.unwrap();
let results = episodic
.search("vegetarian steak eating", 10, false)
.unwrap();
let conflicts = detect_conflicts(&results);
assert_eq!(
conflicts.len(),
1,
"the vegetarian/steak pair must be flagged, got {conflicts:?}"
);
let conflict = &conflicts[0];
assert!(conflict.later_content.contains("steak"));
assert!(conflict.earlier_content.contains("animal products"));
assert!(
conflict.shared_terms.iter().any(|t| t == "vegetarian"),
"shared terms must include the topic: {:?}",
conflict.shared_terms
);
}
#[test]
fn detect_conflicts_ignores_plain_statements_and_assistant_turns() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = store.episodic();
episodic
.append(&user_statement(1, "My favorite coffee is dark roast."))
.unwrap();
episodic
.append(&user_statement(2, "I love coffee with breakfast."))
.unwrap();
episodic
.append(&assistant_response(
3,
"You mentioned you no longer like tea!",
))
.unwrap();
let results = episodic.search("coffee tea breakfast", 10, false).unwrap();
assert!(detect_conflicts(&results).is_empty());
}
#[test]
fn detect_conflicts_skips_identical_content() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = store.episodic();
let record = user_statement(1, "I'm vegetarian now, but I changed my mind.");
let duplicate = user_statement(2, "I'm vegetarian now, but I changed my mind.");
episodic.append(&record).unwrap();
episodic.append(&duplicate).unwrap();
let results = episodic.search("vegetarian", 10, false).unwrap();
assert!(detect_conflicts(&results).is_empty());
}
#[test]
fn append_get_transition_and_reopen_roundtrip() {
let tmp = tempdir().unwrap();
let session = uuid::Uuid::new_v4();
let record = EpisodicRecord::new(
Some(session),
2,
EpisodicKind::Decision,
"use the raw episodic lane",
Provenance::new(ProvenanceSource::User).with_actor("test"),
);
let id = record.id;
{
let store = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = store.episodic();
episodic.append(&record).unwrap();
assert_eq!(episodic.get(id).unwrap().unwrap(), record);
episodic
.transition(
id,
MemoryTransition::Supersede {
replacement: uuid::Uuid::new_v4(),
},
)
.unwrap();
assert!(matches!(
episodic.get(id).unwrap().unwrap().validity,
ValidityState::Superseded { .. }
));
}
let reopened = MemoryStore::open_default(tmp.path()).unwrap();
let records = reopened.episodic().scan(Some(session), 10).unwrap();
assert_eq!(records.len(), 1);
assert_eq!(records[0].id, id);
}
#[test]
fn duplicate_append_is_rejected() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let record = sample_record(1, "once");
store.episodic().append(&record).unwrap();
let error = store.episodic().append(&record).unwrap_err();
assert!(error.to_string().contains("already exists"));
}
#[test]
fn raw_search_returns_canonical_records_and_skips_revoked_by_default() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let active = sample_record(1, "Rust memory retrieval");
let revoked = sample_record(2, "Rust memory retrieval old");
let revoked_id = revoked.id;
store.episodic().append(&active).unwrap();
store.episodic().append(&revoked).unwrap();
store
.episodic()
.transition(
revoked_id,
MemoryTransition::Revoke {
reason: "stale".into(),
},
)
.unwrap();
let current = store
.episodic()
.search("memory retrieval", 10, false)
.unwrap();
assert_eq!(current.len(), 1);
assert_eq!(current[0].record.id, active.id);
let all = store
.episodic()
.search("memory retrieval", 10, true)
.unwrap();
assert_eq!(all.len(), 2);
}
#[test]
fn overlong_terms_index_and_search_without_bad_valsize() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let long_title = "x".repeat(512);
let record = sample_record(1, &format!("session title {long_title}"));
store.episodic().append(&record).unwrap();
assert_eq!(
store.episodic().get(record.id).unwrap(),
Some(record.clone())
);
let hits = store.episodic().search(&long_title, 10, false).unwrap();
assert_eq!(hits.len(), 1, "long-token query must still match");
assert_eq!(hits[0].record.id, record.id);
assert_eq!(index_safe_term(&"a".repeat(480)).len(), 480);
let hashed = index_safe_term(&"a".repeat(481));
assert!(hashed.starts_with("~h:"), "{hashed}");
assert_eq!(hashed.len(), 67);
}
#[test]
fn sidecar_health_exposes_incomplete_derived_index_and_rebuild_heals() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
store
.episodic()
.append(&sample_record(1, "durable raw fact"))
.unwrap();
assert_eq!(store.episodic_sidecar_health().unwrap(), (1, false));
{
let db = store.env().open_db(Some("episodic_terms_v2")).unwrap();
let mut tx = store.env().begin_rw_txn().unwrap();
tx.clear_db(db).unwrap();
tx.commit().unwrap();
}
assert_eq!(
store.episodic_sidecar_health().unwrap(),
(1, true),
"records without postings must be visible as degraded"
);
let rebuilt = store.episodic().rebuild_sidecar().unwrap();
assert_eq!(rebuilt, 1);
assert_eq!(store.episodic_sidecar_health().unwrap(), (1, false));
let hits = store.episodic().search("durable fact", 10, false).unwrap();
assert_eq!(hits.len(), 1);
}
#[test]
fn append_batch_indexes_once_and_preserves_search() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let first = sample_record(1, "Dr. Patel scheduled a follow-up appointment");
let second = sample_record(2, "unrelated grocery list");
let first_id = first.id;
store.episodic().append_batch(&[first, second]).unwrap();
let hits = store
.episodic()
.search("patel appointment", 10, false)
.unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].record.id, first_id);
}
#[test]
fn rejected_append_batch_preserves_prior_state_after_reopen() {
let tmp = tempdir().unwrap();
let original = sample_record(1, "acknowledged original record");
let rejected_new = sample_record(2, "must not survive rejected batch");
{
let store = MemoryStore::open_default(tmp.path()).unwrap();
store.episodic().append(&original).unwrap();
let error = store
.episodic()
.append_batch(&[rejected_new.clone(), original.clone()])
.unwrap_err();
assert!(error.to_string().contains("already exists"));
assert_eq!(
store.episodic().get(original.id).unwrap(),
Some(original.clone())
);
assert_eq!(store.episodic().get(rejected_new.id).unwrap(), None);
}
let reopened = MemoryStore::open_default(tmp.path()).unwrap();
assert_eq!(
reopened.episodic().get(original.id).unwrap(),
Some(original)
);
assert_eq!(reopened.episodic().get(rejected_new.id).unwrap(), None);
}
#[test]
fn lost_episodic_sidecar_rebuilds_from_raw_records_after_reopen() {
let tmp = tempdir().unwrap();
let record = sample_record(1, "sidecar recovery preserves searchable evidence");
{
let store = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = store.episodic();
episodic.append(&record).unwrap();
assert!(!episodic.sidecar_is_empty().unwrap());
let mut tx = store.env().begin_rw_txn().unwrap();
tx.clear_db(episodic.term_db).unwrap();
tx.commit().unwrap();
assert!(episodic.sidecar_is_empty().unwrap());
}
let reopened = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = reopened.episodic();
assert_eq!(episodic.get(record.id).unwrap(), Some(record.clone()));
assert!(!episodic.sidecar_is_empty().unwrap());
let hits = episodic
.search("sidecar searchable evidence", 10, false)
.unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].record.id, record.id);
}
#[test]
fn append_explicit_batch_redacts_and_skips_private() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let public = sample_record(1, "api_key=supersecret rust retrieval");
let private = sample_record(2, "private rust retrieval").with_visibility(true, false);
let public_id = public.id;
store
.episodic()
.append_explicit_batch(&[public, private], EpisodicCapturePolicy::explicit_only())
.unwrap();
let stored = store.episodic().get(public_id).unwrap().unwrap();
assert!(stored.content.contains("<REDACTED>"));
let hits = store
.episodic()
.search("rust retrieval", 10, false)
.unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].record.id, public_id);
}
#[test]
fn typed_keys_retrieve_vocabulary_mismatch() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let dog = sample_record(1, "My Golden Retriever loves the park");
let other = sample_record(2, "I bought a yellow dress");
let dog_id = dog.id;
store.episodic().append_batch(&[dog, other]).unwrap();
let hits = store
.episodic()
.search("What breed is my dog?", 5, false)
.unwrap();
assert_eq!(hits[0].record.id, dog_id);
}
#[test]
fn planner_boosts_temporal_date_match() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let dated = sample_record(1, "I volunteered on February 14th at the animal shelter");
let other = sample_record(2, "I volunteered at the community garden last summer");
let dated_id = dated.id;
store.episodic().append_batch(&[dated, other]).unwrap();
let hits = store
.episodic()
.search("When did I volunteer at the animal shelter?", 5, false)
.unwrap();
assert_eq!(hits[0].record.id, dated_id);
}
#[test]
#[ignore = "manual in-process latency profile"]
fn profile_ingest_and_search_latency() {
fn timed_ms(label: &str, repeats: u32, mut work: impl FnMut()) {
let start = std::time::Instant::now();
for _ in 0..repeats {
work();
}
let elapsed = start.elapsed();
println!(
"{label}: {:.3} ms (n={repeats})",
elapsed.as_secs_f64() * 1000.0 / f64::from(repeats)
);
}
let search_records: Vec<EpisodicRecord> = (0..10_000)
.map(|n| {
sample_record(
n,
if n % 5 == 0 {
"Rust memory retrieval benchmark item"
} else {
"Unrelated episodic record"
},
)
})
.collect();
timed_ms("append_single_1000", 1, || {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
for n in 0..1_000 {
store.episodic().append(&sample_record(n, "once")).unwrap();
}
});
timed_ms("append_batch_1000", 1, || {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let records: Vec<EpisodicRecord> =
(0..1_000).map(|n| sample_record(n, "once")).collect();
store.episodic().append_batch(&records).unwrap();
});
let tmp = tempdir().unwrap();
{
let store = MemoryStore::open_default(tmp.path()).unwrap();
store.episodic().append_batch(&search_records).unwrap();
}
let cold = MemoryStore::open_default(tmp.path()).unwrap();
timed_ms("cold_search_10000", 1, || {
let hits = cold
.episodic()
.search("rust memory retrieval", 10, false)
.unwrap();
assert!(!hits.is_empty());
});
timed_ms("warm_search_10000", 50, || {
let hits = cold
.episodic()
.search("rust memory retrieval", 10, false)
.unwrap();
assert!(!hits.is_empty());
});
}
#[test]
#[ignore = "manual in-process latency profile at realistic scale"]
fn profile_search_latency_25k_realistic() {
const TOTAL: u64 = 25_000;
const SESSIONS: u64 = 50;
let topics = [
"bookshelf",
"guitar",
"vegetarian",
"portfolio",
"commute",
"grandmother",
"chemistry",
"marathon",
"internship",
"yoga",
"spam filter",
"projector",
"swimming",
"cousin",
"bank account",
"book club",
"recipe",
"journal subscription",
"laptop",
"hiking",
];
let fillers = [
"We discussed the plan for the weekend and agreed on the schedule.",
"The meeting notes were circulated and everyone acknowledged them.",
"I explained my reasoning and the group considered the proposal.",
"After the presentation we reviewed the feedback together.",
"She mentioned the deadline and we adjusted the timeline accordingly.",
];
let records: Vec<EpisodicRecord> = (0..TOTAL)
.map(|n| {
let session = n * SESSIONS / TOTAL;
let topic = topics[(n as usize) % topics.len()];
let filler = fillers[(n as usize) % fillers.len()];
let content = format!(
"Session {session} note {n}: my friend Alice mentioned {topic} while {filler}"
);
let session_id = if n % 4 == 0 {
None
} else {
Some(uuid::Uuid::new_v4())
};
EpisodicRecord::new(
session_id,
n,
if n % 3 == 0 {
EpisodicKind::UserStatement
} else {
EpisodicKind::AssistantResponse
},
content,
Provenance::new(ProvenanceSource::User),
)
})
.collect();
let tmp = tempdir().unwrap();
{
let store = MemoryStore::open_default(tmp.path()).unwrap();
let t = std::time::Instant::now();
store.episodic().append_batch(&records).unwrap();
println!("ingest 25k: {:.1} ms", t.elapsed().as_secs_f64() * 1000.0);
}
let cold = MemoryStore::open_default(tmp.path()).unwrap();
let episodic = cold.episodic();
let t = std::time::Instant::now();
let hits = episodic
.search("guitar grandmother recipe", 10, false)
.unwrap();
println!(
"cold_search: {:.1} ms (hits={})",
t.elapsed().as_secs_f64() * 1000.0,
hits.len()
);
let queries: Vec<String> = (0..20)
.map(|i| {
let a = topics[i * 7 % topics.len()];
let b = topics[(i * 7 + 5) % topics.len()];
format!("{a} {b} weekend plan")
})
.collect();
let start = std::time::Instant::now();
for q in &queries {
let hits = episodic.search(q, 10, false).unwrap();
assert!(!hits.is_empty(), "no hits for {q}");
}
println!(
"warm_search varied p50: {:.2} ms/query",
start.elapsed().as_secs_f64() * 1000.0 / queries.len() as f64
);
let t = std::time::Instant::now();
let hits = episodic
.search("zzzterm zzzother zzzthird", 10, false)
.unwrap();
println!(
"no_match_query (full-scan fallback path): {:.1} ms (hits={})",
t.elapsed().as_secs_f64() * 1000.0,
hits.len()
);
}
#[test]
fn enrichment_bridges_vocabulary_gap_for_theater() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
let answer = sample_record(1, "The production I attended was The Glass Menagerie");
let competing = sample_record(2, "I went to a play at the local community theater");
let answer_id = answer.id;
store.episodic().append_batch(&[answer, competing]).unwrap();
let hits = store
.episodic()
.search(
"What play did I attend at the local community theater?",
5,
false,
)
.unwrap();
assert!(hits.iter().any(|h| h.record.id == answer_id));
}
#[test]
fn enrichment_bridges_vocabulary_gap_for_shelter() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
store.set_episodic_enrichment(crate::enrichment::VocabularyEnrichment::with_defaults());
let answer = sample_record(1, "I rescued a dog from the humane society last week");
let other = sample_record(2, "I bought groceries at the store");
let answer_id = answer.id;
store.episodic().append_batch(&[answer, other]).unwrap();
let hits = store
.episodic()
.search("When did I volunteer at the animal shelter?", 5, false)
.unwrap();
assert!(hits.iter().any(|h| h.record.id == answer_id));
}
#[test]
fn session_boost_favors_sessions_with_multiple_matches() {
let tmp = tempdir().unwrap();
let store = MemoryStore::open_default(tmp.path()).unwrap();
let session_a = uuid::Uuid::new_v4();
let session_b = uuid::Uuid::new_v4();
let a1 = EpisodicRecord::new(
Some(session_a),
1,
EpisodicKind::Observation,
"I love hiking in the mountains",
Provenance::new(ProvenanceSource::User),
);
let a2 = EpisodicRecord::new(
Some(session_a),
2,
EpisodicKind::Observation,
"Hiking in the mountains is great exercise",
Provenance::new(ProvenanceSource::User),
);
let b1 = EpisodicRecord::new(
Some(session_b),
1,
EpisodicKind::Observation,
"Hiking is fun",
Provenance::new(ProvenanceSource::User),
);
store.episodic().append_batch(&[a1, a2, b1]).unwrap();
let hits = store
.episodic()
.search("hiking mountains", 10, false)
.unwrap();
let a_ranks: Vec<usize> = hits
.iter()
.enumerate()
.filter(|(_, h)| h.record.session_id == Some(session_a))
.map(|(i, _)| i)
.collect();
let b_rank = hits
.iter()
.position(|h| h.record.session_id == Some(session_b));
if let Some(br) = b_rank {
assert!(a_ranks.iter().all(|&ar| ar < br));
}
}
}