use crate::MemoryId;
use serde::{Deserialize, Serialize};
use wm_core::{CoreError, Galaxy, Result};
use std::path::Path;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use tantivy::{
Index, IndexReader, IndexWriter, ReloadPolicy,
collector::TopDocs,
doc,
query::QueryParser,
schema::{
Field, STORED, STRING, Schema, TantivyDocument, TextFieldIndexing, TextOptions, Value,
},
};
pub const MAX_INDEX_CONTENT_LEN: usize = 8 * 1024;
pub const MIN_PRINTABLE_RATIO: f32 = 0.9;
pub const STOPWORDS: &[&str] = &[
"a",
"about",
"after",
"again",
"all",
"also",
"am",
"an",
"and",
"any",
"are",
"as",
"at",
"be",
"been",
"being",
"before",
"between",
"both",
"but",
"by",
"can",
"could",
"did",
"do",
"does",
"during",
"each",
"few",
"for",
"from",
"further",
"had",
"has",
"have",
"he",
"her",
"here",
"hers",
"herself",
"him",
"himself",
"his",
"how",
"i",
"if",
"in",
"into",
"is",
"it",
"its",
"itself",
"just",
"me",
"might",
"more",
"most",
"my",
"myself",
"no",
"nor",
"not",
"of",
"off",
"on",
"once",
"only",
"or",
"other",
"our",
"ours",
"ourselves",
"out",
"over",
"own",
"same",
"shall",
"she",
"should",
"so",
"some",
"such",
"than",
"that",
"the",
"their",
"theirs",
"them",
"themselves",
"then",
"there",
"these",
"they",
"this",
"those",
"through",
"to",
"too",
"under",
"until",
"up",
"us",
"very",
"was",
"we",
"were",
"what",
"when",
"where",
"which",
"while",
"who",
"whom",
"why",
"will",
"with",
"would",
"you",
"your",
"yours",
"yourself",
"yourselves",
];
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct SearchOptions {
pub limit: usize,
pub galaxy: Option<Galaxy>,
pub min_score: Option<f32>,
pub relative_floor: Option<f32>,
pub relaxed: bool,
}
impl Default for SearchOptions {
fn default() -> Self {
Self {
limit: 20,
galaxy: None,
min_score: None,
relative_floor: None,
relaxed: false,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SearchResult {
pub memory_id: String,
pub galaxy: String,
pub score: f32,
pub normalized_score: f32,
pub content: String,
}
#[derive(Debug, Default)]
pub struct IndexHealth {
pub successes: AtomicU64,
pub failures: AtomicU64,
last_error: Mutex<String>,
}
impl IndexHealth {
fn record_success(&self) {
self.successes.fetch_add(1, Ordering::Relaxed);
}
fn record_failure(&self, err: &str) {
self.failures.fetch_add(1, Ordering::Relaxed);
if let Ok(mut guard) = self.last_error.lock() {
*guard = err.to_string();
}
}
#[must_use]
pub fn snapshot(&self) -> serde_json::Value {
let successes = self.successes.load(Ordering::Relaxed);
let failures = self.failures.load(Ordering::Relaxed);
let last_error = self
.last_error
.lock()
.map(|g| g.clone())
.unwrap_or_default();
let degraded = failures > 0;
serde_json::json!({
"successes": successes,
"failures": failures,
"degraded": degraded,
"last_error": if last_error.is_empty() { serde_json::Value::Null } else { serde_json::Value::String(last_error) },
})
}
}
fn format_writer_lock_error(err: &str, index_path: &Path) -> String {
let is_lock = err.contains("ock") && (err.contains("Busy") || err.contains("lock"));
if is_lock {
format!(
"Tantivy writer: {err} — the search index at {} is locked by another process. \
A running `wm serve` or `wm daemon` on this store holds it; find it with \
`pgrep -af wm` and stop it, or start this server with --readonly.",
index_path.display()
)
} else {
format!("Tantivy writer: {err}")
}
}
pub struct SearchEngine {
index: Index,
reader: IndexReader,
writer: Mutex<Option<IndexWriter>>,
field_id: Field,
field_galaxy: Field,
field_content: Field,
field_tags: Field,
field_timestamp: Field,
health: IndexHealth,
schema_migrated: bool,
}
impl SearchEngine {
fn build_schema() -> (Schema, Field, Field, Field, Field, Field) {
let mut schema_builder = Schema::builder();
let field_id = schema_builder.add_text_field("memory_id", STRING | STORED);
let field_galaxy = schema_builder.add_text_field("galaxy", STRING | STORED);
let stem_indexing = TextFieldIndexing::default()
.set_tokenizer("en_stem")
.set_index_option(tantivy::schema::IndexRecordOption::WithFreqsAndPositions);
let stem_text = TextOptions::default()
.set_indexing_options(stem_indexing.clone())
.set_stored();
let stem_tags = TextOptions::default().set_indexing_options(stem_indexing);
let field_content = schema_builder.add_text_field("content", stem_text);
let field_tags = schema_builder.add_text_field("tags", stem_tags);
let field_timestamp = schema_builder.add_i64_field("timestamp", STORED);
let schema = schema_builder.build();
(
schema,
field_id,
field_galaxy,
field_content,
field_tags,
field_timestamp,
)
}
fn open_index(path: &Path, schema: &Schema, writable: bool) -> Result<(Index, bool)> {
let directory = tantivy::directory::MmapDirectory::open(path)
.map_err(|e| CoreError::Memory(format!("Tantivy open directory: {e}")))?;
if !writable {
let index = Index::open(directory).map_err(|e| {
CoreError::Memory(format!(
"Tantivy readonly open-existing at {}: {e}",
path.display()
))
})?;
if index.schema() != *schema {
return Err(CoreError::Memory(format!(
"Tantivy index at {} was created with an incompatible schema by an \
older version. Run 'wm reindex' (or start 'wm serve' without \
--readonly) to migrate and rebuild it from the canonical store.",
path.display()
)));
}
return Ok((index, false));
}
match Index::open_or_create(directory, schema.clone()) {
Ok(index) => Ok((index, false)),
Err(tantivy::error::TantivyError::SchemaError(_)) => {
let ts = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |d| d.as_millis());
let file_name = path
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("tantivy");
let backup = path.with_file_name(format!("{file_name}.schema-mismatch.{ts}"));
std::fs::rename(path, &backup).map_err(|e| {
CoreError::Memory(format!(
"Tantivy schema migration — rename old index to {}: {e}",
backup.display()
))
})?;
std::fs::create_dir_all(path).map_err(|e| {
CoreError::Memory(format!("Tantivy schema migration — create index dir: {e}"))
})?;
tracing::warn!(
"Tantivy index schema mismatch — old index moved to {}; creating a fresh \
index (rebuild from LMDB will follow)",
backup.display()
);
let directory = tantivy::directory::MmapDirectory::open(path)
.map_err(|e| CoreError::Memory(format!("Tantivy open directory: {e}")))?;
let index = Index::open_or_create(directory, schema.clone())
.map_err(|e| CoreError::Memory(format!("Tantivy open_or_create: {e}")))?;
Ok((index, true))
}
Err(e) => Err(CoreError::Memory(format!("Tantivy open_or_create: {e}"))),
}
}
pub fn open(path: impl AsRef<Path>) -> Result<Self> {
let path = path.as_ref();
let (schema, field_id, field_galaxy, field_content, field_tags, field_timestamp) =
Self::build_schema();
let (index, schema_migrated) = Self::open_index(path, &schema, true)?;
let reader = index
.reader_builder()
.reload_policy(ReloadPolicy::OnCommitWithDelay)
.try_into()
.map_err(|e| CoreError::Memory(format!("Tantivy reader: {e}")))?;
let writer = index
.writer(50_000_000)
.map_err(|e| CoreError::Memory(format_writer_lock_error(&e.to_string(), path)))?;
Ok(Self {
index,
reader,
writer: Mutex::new(Some(writer)),
field_id,
field_galaxy,
field_content,
field_tags,
field_timestamp,
health: IndexHealth::default(),
schema_migrated,
})
}
pub fn open_readonly(path: impl AsRef<Path>) -> Result<Self> {
let path = path.as_ref();
tracing::warn!(
"read-only search index opened at {} — it will not observe writes made \
after this point; restart the read-only server to pick up new memories",
path.display()
);
let (schema, field_id, field_galaxy, field_content, field_tags, field_timestamp) =
Self::build_schema();
let (index, schema_migrated) = Self::open_index(path, &schema, false)?;
let reader = index
.reader_builder()
.reload_policy(ReloadPolicy::OnCommitWithDelay)
.try_into()
.map_err(|e| CoreError::Memory(format!("Tantivy reader: {e}")))?;
Ok(Self {
index,
reader,
writer: Mutex::new(None),
field_id,
field_galaxy,
field_content,
field_tags,
field_timestamp,
health: IndexHealth::default(),
schema_migrated,
})
}
#[must_use]
pub const fn schema_migrated(&self) -> bool {
self.schema_migrated
}
#[must_use]
pub const fn health(&self) -> &IndexHealth {
&self.health
}
pub fn count_docs_in_galaxy(&self, galaxy: &str) -> Result<usize> {
self.reader
.reload()
.map_err(|e| CoreError::Memory(format!("Tantivy reader reload: {e}")))?;
let searcher = self.reader.searcher();
let term = tantivy::Term::from_field_text(self.field_galaxy, galaxy);
let query = tantivy::query::TermQuery::new(term, tantivy::schema::IndexRecordOption::Basic);
let count = searcher
.search(&query, &tantivy::collector::Count)
.map_err(|e| CoreError::Memory(format!("Tantivy count_docs: {e}")))?;
Ok(count)
}
pub fn indexed_ids_in_galaxy(&self, galaxy: &str) -> Result<std::collections::HashSet<String>> {
self.reader
.reload()
.map_err(|e| CoreError::Memory(format!("Tantivy reader reload: {e}")))?;
let searcher = self.reader.searcher();
let term = tantivy::Term::from_field_text(self.field_galaxy, galaxy);
let query = tantivy::query::TermQuery::new(term, tantivy::schema::IndexRecordOption::Basic);
let count = self.count_docs_in_galaxy(galaxy)?;
let hits: std::collections::HashSet<tantivy::DocAddress> = searcher
.search(&query, &tantivy::collector::DocSetCollector)
.map_err(|e| CoreError::Memory(format!("Tantivy indexed_ids: {e}")))?;
let mut out = std::collections::HashSet::with_capacity(count);
for addr in hits {
let doc: TantivyDocument = searcher
.doc(addr)
.map_err(|e| CoreError::Memory(format!("Tantivy get doc: {e}")))?;
if let Some(id) = doc.get_first(self.field_id).and_then(|v| v.as_str()) {
out.insert(id.to_string());
}
}
Ok(out)
}
pub fn is_readonly(&self) -> bool {
self.writer.lock().map_or(true, |g| g.is_none())
}
pub fn writer(&self) -> Result<std::sync::MutexGuard<'_, Option<IndexWriter>>> {
let guard = self
.writer
.lock()
.map_err(|_| CoreError::Memory("Tantivy writer mutex poisoned".into()))?;
if guard.is_none() {
return Err(CoreError::Memory(
"Tantivy writer unavailable: index opened read-only".into(),
));
}
Ok(guard)
}
pub fn add_document(
&self,
writer: &mut Option<IndexWriter>,
memory_id: &str,
galaxy: &str,
content: &str,
tags: &[String],
timestamp: i64,
) -> Result<()> {
let writer = writer.as_mut().ok_or_else(|| {
CoreError::Memory("Tantivy writer unavailable: index opened read-only".into())
})?;
let Some(clean_content) = sanitize_content_for_index(content) else {
tracing::debug!("Skipping index of memory {memory_id}: content failed sanitization");
return Ok(());
};
let tags_str = tags.join(" ");
let doc = doc!(
self.field_id => memory_id,
self.field_galaxy => galaxy,
self.field_content => clean_content,
self.field_tags => tags_str,
self.field_timestamp => timestamp,
);
match writer.add_document(doc) {
Ok(_) => {
self.health.record_success();
Ok(())
}
Err(e) => {
let msg = format!("Tantivy add_document: {e}");
self.health.record_failure(&msg);
Err(CoreError::Memory(msg))
}
}
}
pub fn index_memory(
&self,
writer: &mut Option<IndexWriter>,
mem: &crate::memory::Memory,
) -> Result<()> {
self.add_document(
writer,
&mem.metadata.id.to_string(),
mem.metadata.galaxy.db_name(),
&mem.content,
&mem.metadata.tags,
mem.metadata.created_at.timestamp(),
)
}
pub fn delete_document(&self, writer: &mut Option<IndexWriter>, memory_id: &str) -> Result<()> {
let writer = writer.as_mut().ok_or_else(|| {
CoreError::Memory("Tantivy writer unavailable: index opened read-only".into())
})?;
let term = tantivy::Term::from_field_text(self.field_id, memory_id);
writer.delete_term(term);
Ok(())
}
pub fn delete_by_galaxy(&self, writer: &mut Option<IndexWriter>, galaxy: &str) -> Result<()> {
let writer = writer.as_mut().ok_or_else(|| {
CoreError::Memory("Tantivy writer unavailable: index opened read-only".into())
})?;
let term = tantivy::Term::from_field_text(self.field_galaxy, galaxy);
writer.delete_term(term);
Ok(())
}
pub fn commit(&self, writer: &mut Option<IndexWriter>) -> Result<()> {
let writer = writer.as_mut().ok_or_else(|| {
CoreError::Memory("Tantivy writer unavailable: index opened read-only".into())
})?;
writer
.commit()
.map_err(|e| CoreError::Memory(format!("Tantivy commit: {e}")))?;
self.reader
.reload()
.map_err(|e| CoreError::Memory(format!("Tantivy reload: {e}")))?;
Ok(())
}
pub fn search(&self, query: &str, limit: usize) -> Result<Vec<SearchResult>> {
let opts = SearchOptions {
limit,
..SearchOptions::default()
};
self.search_opt(query, &opts)
}
pub fn search_in_galaxy(
&self,
query: &str,
galaxy: Option<Galaxy>,
limit: usize,
) -> Result<Vec<SearchResult>> {
let opts = SearchOptions {
limit,
galaxy,
..SearchOptions::default()
};
self.search_opt(query, &opts)
}
pub fn search_opt(&self, query: &str, opts: &SearchOptions) -> Result<Vec<SearchResult>> {
let stripped = strip_stopwords(query);
let sanitized = sanitize_tantivy_query(&stripped);
if sanitized.trim().is_empty() {
return Ok(Vec::new());
}
let searcher = self.reader.searcher();
let query_parser =
QueryParser::for_index(&self.index, vec![self.field_content, self.field_tags]);
let parsed = parse_query_with_fallback(&query_parser, &sanitized);
let collector = TopDocs::with_limit(opts.limit).order_by_score();
let top_docs = searcher
.search(&parsed, &collector)
.map_err(|e| CoreError::Memory(format!("Tantivy search: {e}")))?;
let top_score = top_docs.first().map_or(0.0, |(score, _)| *score);
let absolute_floor = opts.min_score.unwrap_or(f32::MIN);
let relative_floor = opts
.relative_floor
.map_or(f32::MIN, |ratio| top_score * ratio);
let query_tokens = query_stem_tokens(&stripped);
let coverage_floor = if query_tokens.len() >= 3 { 2 } else { 1 };
let mut results = Vec::new();
for (score, doc_address) in top_docs {
if score < absolute_floor || score < relative_floor {
continue;
}
let doc: TantivyDocument = searcher
.doc(doc_address)
.map_err(|e| CoreError::Memory(format!("Tantivy get doc: {e}")))?;
let memory_id = doc
.get_first(self.field_id)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let doc_galaxy = doc
.get_first(self.field_galaxy)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
if let Some(g) = opts.galaxy {
if doc_galaxy != g.db_name() {
continue;
}
}
let content = doc
.get_first(self.field_content)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
if coverage_floor > 1 {
let hits = count_token_hits(&content, &stripped);
if hits < coverage_floor {
continue;
}
}
let boosted_score = if query_tokens.is_empty() {
score
} else {
let hits = count_token_hits(&content, &stripped);
let ratio = hits as f32 / query_tokens.len() as f32;
score * 0.1f32.mul_add(ratio, 1.0)
};
results.push(SearchResult {
memory_id,
galaxy: doc_galaxy,
score: boosted_score,
normalized_score: 0.0, content: scrub_text(&content),
});
}
results.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
let top_boosted = results.first().map_or(0.0, |r| r.score);
for r in &mut results {
r.normalized_score = if top_boosted > 0.0 {
r.score / top_boosted
} else {
0.0
};
}
Ok(results)
}
pub fn search_ids(&self, query: &str, limit: usize) -> Result<Vec<MemoryId>> {
let results = self.search(query, limit)?;
Ok(results
.into_iter()
.filter_map(|r| uuid::Uuid::parse_str(&r.memory_id).ok())
.collect())
}
}
fn parse_query_with_fallback(parser: &QueryParser, query: &str) -> Box<dyn tantivy::query::Query> {
match parser.parse_query(query) {
Ok(parsed) => parsed,
Err(_) => parser.parse_query_lenient(query).0,
}
}
#[must_use]
pub fn sanitize_tantivy_query(input: &str) -> String {
if input.trim().is_empty() {
return String::new();
}
input
.split_whitespace()
.filter(|term| term.chars().any(char::is_alphanumeric))
.map(|term| {
if term_needs_quoting(term) {
let escaped = term.replace('\\', "\\\\").replace('"', "\\\"");
format!("\"{escaped}\"")
} else {
term.to_string()
}
})
.collect::<Vec<_>>()
.join(" ")
}
#[must_use]
fn term_needs_quoting(term: &str) -> bool {
if term.starts_with('+') || term.starts_with('-') || term.starts_with('!') {
return true;
}
if term == "AND" || term == "OR" || term == "NOT" {
return true;
}
if term.contains("&&") || term.contains("||") {
return true;
}
term.chars().any(|c| {
matches!(
c,
'(' | ')' | '{' | '}' | '[' | ']' | '^' | '"' | '~' | '*' | '?' | ':' | '\\' | '/'
)
})
}
#[must_use]
pub fn strip_stopwords(query: &str) -> String {
query
.split_whitespace()
.filter(|term| !STOPWORDS.contains(&term.to_lowercase().as_str()))
.collect::<Vec<_>>()
.join(" ")
}
#[must_use]
fn query_stem_tokens(stripped_query: &str) -> Vec<String> {
stem_tokens(stripped_query)
}
#[must_use]
fn stem_tokens(text: &str) -> Vec<String> {
let mut tokens: Vec<String> = Vec::new();
for term in text
.split(|c: char| !c.is_alphanumeric())
.filter(|term| term.len() > 1)
{
let stemmed = simple_stem(&term.to_lowercase());
if !tokens.contains(&stemmed) {
tokens.push(stemmed);
}
}
tokens
}
#[must_use]
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()
}
#[must_use]
fn count_token_hits(content: &str, stripped_query: &str) -> usize {
let query_tokens = query_stem_tokens(stripped_query);
if query_tokens.is_empty() {
return 0;
}
let content_stems: std::collections::HashSet<String> =
stem_tokens(content).into_iter().collect();
query_tokens
.iter()
.filter(|t| content_stems.contains(*t))
.count()
}
#[must_use]
pub fn sanitize_content_for_index(content: &str) -> Option<String> {
if content.trim().is_empty() {
return None;
}
if content.as_bytes().contains(&0) {
return None;
}
let total = content.chars().count();
if total == 0 {
return None;
}
let printable = content.chars().filter(|c| !c.is_control()).count();
if (printable as f32 / total as f32) < MIN_PRINTABLE_RATIO {
return None;
}
let cleaned = scrub_text(content);
let capped: String = cleaned.chars().take(MAX_INDEX_CONTENT_LEN).collect();
if capped.trim().is_empty() {
None
} else {
Some(capped)
}
}
#[must_use]
pub fn scrub_text(content: &str) -> String {
let mut out = String::with_capacity(content.len().min(MAX_INDEX_CONTENT_LEN));
for c in content.chars().take(MAX_INDEX_CONTENT_LEN) {
if c.is_control() && c != '\n' && c != '\t' && c != '\r' {
out.push(' ');
} else {
out.push(c);
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
fn open_engine() -> (tempfile::TempDir, SearchEngine) {
let tmp = tempdir().unwrap();
let engine = SearchEngine::open(tmp.path()).unwrap();
(tmp, engine)
}
fn write_incompatible_index(dir: &Path) {
std::fs::create_dir_all(dir).unwrap();
let mut builder = Schema::builder();
builder.add_text_field("legacy", STRING | STORED);
let schema = builder.build();
let directory = tantivy::directory::MmapDirectory::open(dir).unwrap();
Index::open_or_create(directory, schema).unwrap();
}
#[test]
fn open_migrates_incompatible_schema() {
let tmp = tempdir().unwrap();
let dir = tmp.path().join("tantivy");
write_incompatible_index(&dir);
let engine = SearchEngine::open(&dir).unwrap();
assert!(
engine.schema_migrated(),
"incompatible schema must trigger migration"
);
let backups: Vec<_> = std::fs::read_dir(tmp.path())
.unwrap()
.filter_map(std::result::Result::ok)
.filter(|e| e.file_name().to_string_lossy().contains("schema-mismatch"))
.collect();
assert_eq!(backups.len(), 1, "old index must be backed up exactly once");
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"33333333-3333-3333-3333-333333333333",
"codex",
"fresh index after migration",
&[],
1700000000,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("fresh index", 10).unwrap();
assert_eq!(results.len(), 1);
}
#[test]
fn writer_lock_error_names_path_and_hint() {
let err = format_writer_lock_error(
"Failed to acquire Lockfile: LockBusy. Some(\"...\")",
Path::new("/store/x/tantivy"),
);
assert!(
err.contains("/store/x/tantivy"),
"must name the index path: {err}"
);
assert!(
err.contains("pgrep -af wm"),
"must include the diagnostic hint: {err}"
);
assert!(
err.contains("--readonly"),
"must offer the readonly alternative: {err}"
);
let other = format_writer_lock_error("disk full", Path::new("/s/t"));
assert!(other.starts_with("Tantivy writer: disk full"));
assert!(!other.contains("pgrep"));
}
#[test]
fn open_readonly_rejects_incompatible_schema() {
let tmp = tempdir().unwrap();
let dir = tmp.path().join("tantivy");
write_incompatible_index(&dir);
let err = match SearchEngine::open_readonly(&dir) {
Ok(_) => panic!("read-only open must reject an incompatible schema"),
Err(e) => e,
};
assert!(
format!("{err}").contains("wm reindex"),
"read-only mismatch must point at wm reindex, got: {err}"
);
let siblings: Vec<_> = std::fs::read_dir(tmp.path())
.unwrap()
.filter_map(std::result::Result::ok)
.filter(|e| e.file_name().to_string_lossy().contains("schema-mismatch"))
.collect();
assert!(siblings.is_empty(), "read-only open must not migrate");
}
#[test]
fn open_readonly_rejects_existing_empty_directory_without_creating_files() {
let tmp = tempdir().unwrap();
let dir = tmp.path().join("tantivy");
std::fs::create_dir_all(&dir).unwrap();
let err = match SearchEngine::open_readonly(&dir) {
Ok(_) => panic!("readonly open unexpectedly initialized an empty index"),
Err(err) => err,
};
assert!(format!("{err}").contains("readonly open-existing"));
assert!(
std::fs::read_dir(&dir).unwrap().next().is_none(),
"readonly open must not materialize Tantivy metadata or segments"
);
}
#[test]
fn reopen_matching_schema_not_migrated() {
let tmp = tempdir().unwrap();
let dir = tmp.path().join("tantivy");
std::fs::create_dir_all(&dir).unwrap();
let first = SearchEngine::open(&dir).unwrap();
assert!(!first.schema_migrated());
drop(first);
let second = SearchEngine::open(&dir).unwrap();
assert!(
!second.schema_migrated(),
"matching schema must not migrate"
);
drop(second);
let third = SearchEngine::open_readonly(&dir).unwrap();
assert!(!third.schema_migrated());
}
#[test]
fn index_and_search_basic() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"The Rust programming language is fast and safe",
&["rust".into(), "programming".into()],
1700000000,
)
.unwrap();
engine
.add_document(
&mut writer,
"22222222-2222-2222-2222-222222222222",
"codex",
"Python is great for data science",
&["python".into(), "data".into()],
1700000001,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("rust", 10).unwrap();
assert!(!results.is_empty());
assert_eq!(results[0].memory_id, "11111111-1111-1111-1111-111111111111");
}
#[test]
fn search_by_tag() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"memory about systems",
&["rust".into()],
1700000000,
)
.unwrap();
engine
.add_document(
&mut writer,
"22222222-2222-2222-2222-222222222222",
"codex",
"memory about cooking",
&["food".into()],
1700000001,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("rust", 10).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].memory_id, "11111111-1111-1111-1111-111111111111");
}
#[test]
fn search_filtered_by_galaxy() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"important knowledge",
&[],
1700000000,
)
.unwrap();
engine
.add_document(
&mut writer,
"22222222-2222-2222-2222-222222222222",
"research",
"important findings",
&[],
1700000001,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine
.search_in_galaxy("important", Some(Galaxy::Codex), 10)
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].galaxy, "codex");
}
#[test]
fn delete_document_from_index() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"deletable content",
&[],
1700000000,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("deletable", 10).unwrap();
assert_eq!(results.len(), 1);
engine
.delete_document(&mut writer, "11111111-1111-1111-1111-111111111111")
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("deletable", 10).unwrap();
assert_eq!(results.len(), 0);
}
#[test]
fn search_empty_index() {
let (_tmp, engine) = open_engine();
let results = engine.search("anything", 10).unwrap();
assert!(results.is_empty());
}
#[test]
fn search_ids_returns_uuids() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"unique content about rust",
&[],
1700000000,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let ids = engine.search_ids("rust", 10).unwrap();
assert_eq!(ids.len(), 1);
assert_eq!(
ids[0],
uuid::Uuid::parse_str("11111111-1111-1111-1111-111111111111").unwrap()
);
}
#[test]
fn sanitize_leaves_plain_terms_unquoted() {
let result = sanitize_tantivy_query("hello world");
assert_eq!(result, "hello world");
}
#[test]
fn sanitize_drops_punct_only_terms() {
let result = sanitize_tantivy_query("*");
assert_eq!(result, "");
}
#[test]
fn sanitize_escapes_boolean_operators() {
let result = sanitize_tantivy_query("NOT secret");
assert_eq!(result, "\"NOT\" secret");
}
#[test]
fn sanitize_escapes_field_syntax() {
let result = sanitize_tantivy_query("content:secret");
assert_eq!(result, "\"content:secret\"");
}
#[test]
fn sanitize_escapes_quotes() {
let result = sanitize_tantivy_query("test\"injection");
assert!(
result.contains("\\\""),
"embedded quotes should be escaped: {result}"
);
}
#[test]
fn sanitize_escapes_trailing_backslash_token() {
let result = sanitize_tantivy_query("C:\\Users\\temp\\");
assert_eq!(
result, "\"C:\\\\Users\\\\temp\\\\\"",
"backslashes must be doubled inside quoted terms"
);
}
#[test]
fn lenient_fallback_never_fails_on_malformed_input() {
let (_tmp, engine) = open_engine();
let parser =
QueryParser::for_index(&engine.index, vec![engine.field_content, engine.field_tags]);
let searcher = engine.reader.searcher();
let collector = TopDocs::with_limit(1).order_by_score();
for malformed in ["\"unterminated", "field:(\"", "\\", "AND NOT OR"] {
let parsed = parse_query_with_fallback(&parser, malformed);
searcher
.search(&parsed, &collector)
.unwrap_or_else(|e| panic!("lenient query {malformed:?} must execute: {e}"));
}
}
#[test]
fn sanitize_empty_returns_empty() {
assert_eq!(sanitize_tantivy_query(""), "");
assert_eq!(sanitize_tantivy_query(" "), "");
}
#[test]
fn sanitize_preserves_alphanumeric() {
let result = sanitize_tantivy_query("rust programming 2024");
assert_eq!(result, "rust programming 2024");
}
#[test]
fn sanitize_preserves_hyphenated_compounds() {
let result = sanitize_tantivy_query("antigravity antigravity-project-test");
assert_eq!(result, "antigravity antigravity-project-test");
}
#[test]
fn strip_stopwords_removes_common_words() {
assert_eq!(
strip_stopwords("smoke test from wmClient"),
"smoke test wmClient"
);
assert_eq!(strip_stopwords("the from and or"), "");
assert_eq!(strip_stopwords("Rust ownership"), "Rust ownership");
assert_eq!(strip_stopwords(""), "");
}
#[test]
fn strip_stopwords_is_case_insensitive() {
assert_eq!(strip_stopwords("FROM The And"), "");
}
#[test]
fn sanitize_content_skips_null_bytes() {
let content = "binary\x00garbage\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00";
assert!(sanitize_content_for_index(content).is_none());
}
#[test]
fn sanitize_content_skips_low_printable_ratio() {
let content = "\u{01}\u{02}\u{03}\u{04}\u{05}hello";
assert!(sanitize_content_for_index(content).is_none());
}
#[test]
fn sanitize_content_scrubs_and_caps() {
let content = "clean text\u{01}with one control char";
let cleaned = sanitize_content_for_index(content).unwrap();
assert!(!cleaned.contains('\u{01}'));
assert!(cleaned.starts_with("clean text with one control char"));
let long = "a".repeat(MAX_INDEX_CONTENT_LEN + 1000);
let capped = sanitize_content_for_index(&long).unwrap();
assert_eq!(capped.chars().count(), MAX_INDEX_CONTENT_LEN);
}
#[test]
fn sanitize_content_skips_empty() {
assert!(sanitize_content_for_index("").is_none());
assert!(sanitize_content_for_index(" ").is_none());
}
#[test]
fn scrub_text_replaces_control_chars() {
let result = scrub_text("a\u{01}b\nc\td\u{7f}e");
assert_eq!(result, "a b\nc\td e");
}
#[test]
fn add_document_skips_binary_content() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"\u{00}\u{01}\u{02}raw serialized bytes",
&[],
1700000000,
)
.unwrap();
engine
.add_document(
&mut writer,
"22222222-2222-2222-2222-222222222222",
"codex",
"clean searchable text",
&[],
1700000001,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("serialized", 10).unwrap();
assert!(results.is_empty(), "binary content must not be indexed");
let results = engine.search("clean", 10).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].memory_id, "22222222-2222-2222-2222-222222222222");
}
fn index_alpha_pair(engine: &SearchEngine) {
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"alpha",
&[],
1700000000,
)
.unwrap();
let filler = format!("alpha {}", "zzz ".repeat(400));
engine
.add_document(
&mut writer,
"22222222-2222-2222-2222-222222222222",
"codex",
&filler,
&[],
1700000001,
)
.unwrap();
engine.commit(&mut writer).unwrap();
}
#[test]
fn search_absolute_min_score_filters_weak_matches() {
let (_tmp, engine) = open_engine();
index_alpha_pair(&engine);
let base = engine.search("alpha", 10).unwrap();
assert_eq!(base.len(), 2);
let (hi, lo) = if base[0].score >= base[1].score {
(base[0].score, base[1].score)
} else {
(base[1].score, base[0].score)
};
assert!(
hi > lo,
"short doc should outscore long doc (hi={hi}, lo={lo})"
);
let mid = f32::midpoint(hi, lo);
let opts = SearchOptions {
limit: 10,
min_score: Some(mid),
..SearchOptions::default()
};
let filtered = engine.search_opt("alpha", &opts).unwrap();
assert_eq!(filtered.len(), 1);
assert!((filtered[0].score - hi).abs() < 1e-3);
}
#[test]
fn search_relative_floor_filters_weak_matches() {
let (_tmp, engine) = open_engine();
index_alpha_pair(&engine);
let opts = SearchOptions {
limit: 10,
relative_floor: Some(0.5),
..SearchOptions::default()
};
let filtered = engine.search_opt("alpha", &opts).unwrap();
assert_eq!(filtered.len(), 1, "weak match must fall below 50% of top");
assert_eq!(
filtered[0].memory_id,
"11111111-1111-1111-1111-111111111111"
);
assert!((filtered[0].normalized_score - 1.0).abs() < 1e-3);
}
#[test]
fn search_all_results_normalized() {
let (_tmp, engine) = open_engine();
index_alpha_pair(&engine);
let results = engine.search("alpha", 10).unwrap();
assert_eq!(results.len(), 2);
assert!((results[0].normalized_score - 1.0).abs() < 1e-3);
for r in &results[1..] {
assert!(r.normalized_score <= 1.0);
assert!(r.normalized_score > 0.0);
}
}
#[test]
fn search_stemming_matches_morphological_variants() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"I graduated with a degree in Business Administration",
&[],
1700000000,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("graduate", 10).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].memory_id, "11111111-1111-1111-1111-111111111111");
let results = engine.search("degrees", 10).unwrap();
assert_eq!(results.len(), 1);
}
#[test]
fn search_incident_query_returns_only_relevant() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
let smoke_id = "11111111-1111-1111-1111-111111111111";
engine
.add_document(
&mut writer,
smoke_id,
"codex",
"smoke test from wmClient: verify recall works",
&[],
1700000000,
)
.unwrap();
let unrelated = [
"NES Evolution and Impact: a history of the console wars",
"Insights on The Gateless Gate: koans and zen practice",
"What the tweet is really saying: a thread analysis",
];
for (i, content) in (1i64..).zip(unrelated.iter()) {
engine
.add_document(
&mut writer,
&format!("22222222-2222-2222-2222-2222222222{i:02}"),
"codex",
content,
&[],
1700000000 + i,
)
.unwrap();
}
engine.commit(&mut writer).unwrap();
let results = engine.search("smoke test from wmClient", 20).unwrap();
assert_eq!(
results.len(),
1,
"only the smoke memory should match: {results:?}"
);
assert_eq!(results[0].memory_id, smoke_id);
assert!(results[0].content.contains("smoke test"));
}
#[test]
fn search_project_compound_query() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"[antigravity:antigravity-project-test]\nQ: how does it work?\nA: details here",
&["project_antigravity-project-test".into()],
1700000000,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine
.search("antigravity antigravity-project-test", 10)
.unwrap();
assert!(
!results.is_empty(),
"project compound query must match the antigravity memory"
);
assert_eq!(results[0].memory_id, "11111111-1111-1111-1111-111111111111");
}
#[test]
fn or_default_filters_partial_matches_via_coverage() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"alpha beta gamma delta",
&[],
1700000000,
)
.unwrap();
engine
.add_document(
&mut writer,
"22222222-2222-2222-2222-222222222222",
"codex",
"alpha only here",
&[],
1700000001,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("alpha beta gamma", 10).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].memory_id, "11111111-1111-1111-1111-111111111111");
let results = engine.search("alpha beta", 10).unwrap();
assert_eq!(results.len(), 2);
}
#[test]
fn coverage_is_case_insensitive() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"Smoke Test for wmClient integration",
&[],
1700000000,
)
.unwrap();
engine
.add_document(
&mut writer,
"22222222-2222-2222-2222-222222222222",
"codex",
"test only here",
&[],
1700000001,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("Smoke Test from wmClient", 10).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].memory_id, "11111111-1111-1111-1111-111111111111");
}
#[test]
fn coverage_matches_stemmed_variants() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"I graduated with a degree in Business Administration",
&[],
1700000000,
)
.unwrap();
engine
.add_document(
&mut writer,
"22222222-2222-2222-2222-222222222222",
"codex",
"degree only here",
&[],
1700000001,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("graduate degree", 10).unwrap();
assert_eq!(results.len(), 2);
assert_eq!(results[0].memory_id, "11111111-1111-1111-1111-111111111111");
}
#[test]
fn coverage_normalizes_possessives_and_punctuation() {
let query = "buy sister's birthday gift";
assert_eq!(
query_stem_tokens(query),
["buy", "sister", "birthday", "gift"]
);
assert_eq!(
count_token_hits("I bought a dress for my sister birthday", query),
2
);
}
#[test]
fn search_stopword_only_query_returns_nothing() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"11111111-1111-1111-1111-111111111111",
"codex",
"some ordinary text",
&[],
1700000000,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("the from and or", 10).unwrap();
assert!(results.is_empty());
}
#[test]
fn wildcard_query_doesnt_match_all() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(&mut writer, "uuid-1", "codex", "first document", &[], 1000)
.unwrap();
engine
.add_document(&mut writer, "uuid-2", "codex", "second document", &[], 2000)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("*", 10).unwrap();
assert!(
results.is_empty(),
"wildcard query should not match all documents after sanitization"
);
}
#[test]
fn field_syntax_query_doesnt_access_other_fields() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(&mut writer, "uuid-1", "secret", "public content", &[], 1000)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("galaxy:secret", 10).unwrap();
assert!(
results.is_empty(),
"field syntax injection should not access non-searchable fields"
);
}
#[test]
fn boolean_operator_doesnt_bypass_search() {
let (_tmp, engine) = open_engine();
let mut writer = engine.writer().unwrap();
engine
.add_document(
&mut writer,
"uuid-1",
"codex",
"important secret data",
&[],
1000,
)
.unwrap();
engine.commit(&mut writer).unwrap();
let results = engine.search("AND secret", 10).unwrap();
assert_eq!(results.len(), 1);
let results = engine.search("AND OR NOT", 10).unwrap();
assert!(
results.is_empty(),
"operator-only query must not bypass search"
);
}
}