Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
132 changes: 123 additions & 9 deletions src/memory/retrieval.rs
Original file line number Diff line number Diff line change
Expand Up @@ -126,8 +126,14 @@ impl<'a> FactRetriever<'a> {

let mut results = Vec::with_capacity(candidates.len());
for fact in candidates {
let fts_score = fts_scores.get(&fact.fact_id).copied().unwrap_or(0.0);
let jaccard_score = jaccard(&query_tokens, &fact_search_tokens(&fact));
let fact_tokens = fact_search_tokens(&fact);
let coverage = term_coverage(&query_tokens, &fact_tokens);
// Weight raw BM25 by distinct-term coverage: a document matching
// every query term must not lose to a shorter document matching
// one term on length normalization alone.
let fts_score =
fts_scores.get(&fact.fact_id).copied().unwrap_or(0.0) * (0.5 + 0.5 * coverage);
let jaccard_score = jaccard(&query_tokens, &fact_tokens);
let holographic_score =
self.holographic_score_with(query, &fact, candidate_vectors.get(&fact.fact_id));
let trust_score = fact.trust_score;
Expand All @@ -149,7 +155,7 @@ impl<'a> FactRetriever<'a> {
holographic_score,
trust_score,
why: Some(format!(
"fts={fts_score:.3}, jaccard={jaccard_score:.3}, holographic={holographic_score:.3}, trust={trust_score:.3}, temporal_decay={temporal_decay:.3}, retrieval_count={retrieval_count}"
"fts={fts_score:.3}, coverage={coverage:.3}, jaccard={jaccard_score:.3}, holographic={holographic_score:.3}, trust={trust_score:.3}, temporal_decay={temporal_decay:.3}, retrieval_count={retrieval_count}"
)),
});
}
Expand Down Expand Up @@ -415,7 +421,7 @@ impl<'a> FactRetriever<'a> {
}
.map_err(|e| db_error("fts_candidates", e))?;

let mut scores = HashMap::new();
let mut ranked = Vec::new();
while let Some(row) = rows
.next()
.await
Expand All @@ -424,13 +430,13 @@ impl<'a> FactRetriever<'a> {
let rank = row
.get::<f64>(1)
.map_err(|e| db_error("fts_candidates", e))?;
scores.insert(
ranked.push((
row.get::<i64>(0)
.map_err(|e| db_error("fts_candidates", e))?,
1.0 / (1.0 + rank.abs()),
);
rank,
));
}
Ok(scores)
Ok(normalize_fts5_ranks(ranked))
}

async fn entity_candidates(
Expand Down Expand Up @@ -649,12 +655,77 @@ fn build_fts_query(query: &str) -> Option<String> {
Some(
tokens
.into_iter()
.map(|token| format!("\"{}\"", token.replace('"', "\"\"")))
.map(|token| {
let quoted = format!("\"{}\"", token.replace('"', "\"\""));
// Prefix-match longer terms so simple morphology reaches the
// index ("install" finds "installing"); short tokens stay
// exact to avoid over-matching.
if token.chars().count() >= PREFIX_MATCH_MIN_CHARS {
format!("{quoted}*")
} else {
quoted
}
})
.collect::<Vec<_>>()
.join(" OR "),
)
}

/// Minimum token length for prefix-form FTS terms and coverage prefix
/// matching. Below this, prefixes over-match ("in" would swallow "install").
const PREFIX_MATCH_MIN_CHARS: usize = 4;

/// Fraction of distinct query tokens the fact's tokens cover, using the same
/// exact-or-prefix semantics as [`build_fts_query`]. BM25 length
/// normalization can rank a short one-term document above a longer document
/// matching every query term; weighting the fts component by coverage keeps
/// multi-term relevance ahead of single-term brevity.
fn term_coverage(query_tokens: &[String], fact_tokens: &[String]) -> f64 {
if query_tokens.is_empty() {
return 0.0;
}
let matched = query_tokens
.iter()
.filter(|q| {
fact_tokens.iter().any(|t| {
t == *q
|| (q.chars().count() >= PREFIX_MATCH_MIN_CHARS && t.starts_with(q.as_str()))
})
})
.count();
matched as f64 / query_tokens.len() as f64
}

fn normalize_fts5_ranks(ranked: Vec<(i64, f64)>) -> HashMap<i64, f64> {
let max_relevance = ranked
.iter()
.map(|(_, rank)| fts5_rank_relevance(*rank))
.fold(0.0_f64, f64::max);
if max_relevance <= f64::EPSILON {
return ranked
.into_iter()
.map(|(fact_id, _)| (fact_id, 0.0))
.collect();
}
ranked
.into_iter()
.map(|(fact_id, rank)| {
(
fact_id,
(fts5_rank_relevance(rank) / max_relevance).clamp(0.0, 1.0),
)
})
.collect()
}

fn fts5_rank_relevance(rank: f64) -> f64 {
if rank.is_finite() {
(-rank).max(0.0)
} else {
0.0
}
}

fn tokenize(text: &str) -> Vec<String> {
let mut tokens = Vec::new();
let mut current = String::new();
Expand Down Expand Up @@ -801,6 +872,49 @@ fn db_error(operation: &str, error: impl fmt::Display) -> TraceDecayError {
mod tests {
use super::*;

#[test]
fn build_fts_query_prefixes_long_tokens_only() {
let q = build_fts_query("install dependencies in db").unwrap();
assert!(q.contains("\"install\"*"), "{q}");
assert!(q.contains("\"dependencies\"*"), "{q}");
// Short tokens stay exact so prefixes cannot over-match.
assert!(!q.contains("\"db\"*"), "{q}");
}

#[test]
fn term_coverage_counts_distinct_terms_with_prefix_reach() {
let query = vec!["install".to_string(), "dependencies".to_string()];
let both = vec!["installing".to_string(), "dependencies".to_string()];
let one = vec!["install".to_string(), "packages".to_string()];
assert!((term_coverage(&query, &both) - 1.0).abs() < f64::EPSILON);
assert!((term_coverage(&query, &one) - 0.5).abs() < f64::EPSILON);
// Short query tokens never prefix-match.
let short = vec!["db".to_string()];
let fact = vec!["dbx".to_string()];
assert!(term_coverage(&short, &fact).abs() < f64::EPSILON);
}

#[test]
fn fts5_ranks_are_normalized_without_losing_small_matches() {
let scores = normalize_fts5_ranks(vec![
(1, -0.000_002),
(2, -0.000_001),
(3, 0.0),
(4, f64::NAN),
]);

assert!((scores[&1] - 1.0).abs() < f64::EPSILON);
assert!((scores[&2] - 0.5).abs() < f64::EPSILON);
assert!(scores[&3].abs() < f64::EPSILON);
assert!(scores[&4].abs() < f64::EPSILON);
}

#[test]
fn lone_common_term_match_keeps_full_fts_weight() {
let scores = normalize_fts5_ranks(vec![(1, -0.000_001)]);
assert!((scores[&1] - 1.0).abs() < f64::EPSILON);
}

#[test]
fn retrieval_reinforcement_boosts_frequently_retrieved_facts() {
let fts = 0.5;
Expand Down
146 changes: 146 additions & 0 deletions tests/mcp_suite/mcp_handler_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3002,6 +3002,152 @@ async fn context_includes_matching_memory_facts() {
);
}

#[tokio::test]
async fn fact_search_ranks_exact_operational_evidence_and_tracks_once() {
let (cg, _env, _dir) = setup_empty_project().await;
let exact = "22 long-lived tracedecay serve processes spanning 0.0.38 through 0.0.47; four 0.0.45 processes hold selected tracedecay.db file descriptors; doctor/upgrade should report stale PIDs/versions/open holders, never kill.";
let unrelated = [
"TraceDecay V2 multi-agent task execution spans several repositories and decomposes into independently claimable task subgraphs with versioned compact context packets.",
"TraceDecay V2 task-graph scoping uses one profile-owned canonical task graph with Kanban, DAG, timeline, workload, initiative, and saved-query projections.",
"TraceDecay V2 task execution relates tickets to threads, sessions, turns, agents, tool calls, files, symbols, worktrees, commits, pull requests, and evidence.",
"TraceDecay V2 may run a daemon-side context scout that observes bounded turn events and emits compact relevance-scored suggestion envelopes.",
"TraceDecay V2 session and LCM retrieval distinguishes current truth from historical evidence and ranks explicit scope, thread, project, worktree, trust, and current-state signals.",
];

let mut contents = vec![exact];
contents.extend(unrelated);
let mut exact_fact_id = None;
for content in &contents {
let added = handle_tool_call(
&cg,
"tracedecay_fact_store",
json!({
"action": "add",
"format": "json",
"content": content,
"category": "decision",
"trust": 0.99,
"source": "fact-ranking-regression"
}),
None,
None,
)
.await
.unwrap();
let added: Value = serde_json::from_str(extract_text(&added.value)).unwrap();
if *content == exact {
exact_fact_id = added["fact"]["fact_id"].as_i64();
}
}
let exact_fact_id = exact_fact_id.expect("exact operational fact should be stored");

let first = handle_tool_call(
&cg,
"tracedecay_fact_store",
json!({
"action": "search",
"format": "json",
"query": "stale tracedecay serve processes versions open database file descriptors doctor upgrade",
"limit": 10,
"min_trust": 0.0
}),
None,
None,
)
.await
.unwrap();
let first: Value = serde_json::from_str(extract_text(&first.value)).unwrap();
let first_results = first["facts"].as_array().expect("fact search results");
assert_eq!(
first_results[0]["fact"]["fact_id"].as_i64(),
Some(exact_fact_id),
"exact operational evidence must outrank unrelated V2 facts: {first}"
);
let after_first = cg.get_fact(exact_fact_id).await.unwrap().unwrap();
assert_eq!(after_first.retrieval_count, 1);
assert_eq!(after_first.access_count, 1);
assert!(after_first.last_retrieved_at.is_some());
assert!(after_first.last_recalled_at.is_some());

let context = handle_tool_call(
&cg,
"tracedecay_context",
json!({
"task": "stale tracedecay serve processes versions open database file descriptors doctor upgrade",
"format": "json",
"memory_limit": 10,
"memory_min_trust": 0.0
}),
None,
None,
)
.await
.unwrap();
let context: Value = serde_json::from_str(extract_text(&context.value)).unwrap();
assert!(context["memory_matches"].as_array().is_some_and(|matches| {
matches
.iter()
.any(|hit| hit["fact"]["fact_id"].as_i64() == Some(exact_fact_id))
}));
let after_context = cg.get_fact(exact_fact_id).await.unwrap().unwrap();
assert_eq!(after_context.retrieval_count, 1);
assert_eq!(after_context.access_count, 1);

let rare = handle_tool_call(
&cg,
"tracedecay_fact_store",
json!({
"action": "search",
"format": "json",
"query": "22 long-lived 0.0.38 0.0.47 four 0.0.45",
"limit": 10,
"min_trust": 0.0
}),
None,
None,
)
.await
.unwrap();
let rare: Value = serde_json::from_str(extract_text(&rare.value)).unwrap();
let rare_results = rare["facts"].as_array().expect("rare-term results");
assert_eq!(
rare_results.len(),
1,
"rare terms should exclude unrelated facts: {rare}"
);
assert_eq!(
rare_results[0]["fact"]["fact_id"].as_i64(),
Some(exact_fact_id)
);
assert!(rare_results[0]["fts_score"].as_f64().unwrap_or_default() > 0.0);
let after_rare = cg.get_fact(exact_fact_id).await.unwrap().unwrap();
assert_eq!(after_rare.retrieval_count, 2);
assert_eq!(after_rare.access_count, 2);

let analytics = handle_tool_call(
&cg,
"tracedecay_analytics",
json!({"section": "facts", "format": "json"}),
None,
None,
)
.await
.unwrap();
let analytics: Value = serde_json::from_str(extract_text(&analytics.value)).unwrap();
assert_eq!(
analytics["facts"]["facts"].as_i64(),
Some(contents.len() as i64)
);
assert_eq!(
analytics["facts"]["retrievals"].as_i64(),
Some(first_results.len() as i64 + rare_results.len() as i64)
);
assert_eq!(
analytics["facts"]["facts_retrieved"].as_i64(),
Some(contents.len() as i64)
);
}

#[tokio::test]
async fn context_memory_controls_filter_disable_and_preserve_markdown() {
let (cg, _dir) = setup_project().await;
Expand Down
Loading