#[path = "delivery.rs"]
pub mod delivery;
use std::cmp::Ordering;
use std::collections::HashMap;
use kimetsu_core::config::{BrokerWeights, StageWeights};
use kimetsu_core::memory::MemoryScope;
use kimetsu_core::{KimetsuResult, ids::new_id};
use rusqlite::{Connection, OptionalExtension, params};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TaskKind {
#[default]
Feature,
Debug,
Refactor,
Docs,
Investigation,
}
pub fn classify_task(task: &str) -> TaskKind {
let lower = task.to_ascii_lowercase();
const DEBUG_KW: &[&str] = &[
"fix",
"bug",
"error",
"fail",
"crash",
"panic",
"regression",
"broken",
"debug",
"stack trace",
"exception",
];
if DEBUG_KW.iter().any(|kw| lower.contains(kw)) {
return TaskKind::Debug;
}
const INVESTIGATE_KW: &[&str] = &[
"investigate",
"analyze",
"understand",
" why ",
"explore",
"find out",
"root cause",
"audit",
"trace",
];
if INVESTIGATE_KW.iter().any(|kw| lower.contains(kw)) {
return TaskKind::Investigation;
}
const REFACTOR_KW: &[&str] = &[
"refactor",
"rename",
"cleanup",
"clean up",
"restructure",
"simplify",
"extract",
"deduplicate",
"reorganize",
];
if REFACTOR_KW.iter().any(|kw| lower.contains(kw)) {
return TaskKind::Refactor;
}
const DOCS_KW: &[&str] = &[
"document",
"readme",
"changelog",
"comment",
"docstring",
"docs",
"tutorial",
"guide",
];
if DOCS_KW.iter().any(|kw| lower.contains(kw)) {
return TaskKind::Docs;
}
TaskKind::Feature
}
fn weights_for_task_kind(base: StageWeights, kind: TaskKind) -> StageWeights {
match kind {
TaskKind::Feature => base,
TaskKind::Debug => renorm(StageWeights {
freshness: base.freshness * 1.6,
..base
}),
TaskKind::Refactor => renorm(StageWeights {
scope: base.scope * 1.6,
..base
}),
TaskKind::Investigation => renorm(StageWeights {
relevance: base.relevance * 1.4,
..base
}),
TaskKind::Docs => renorm(StageWeights {
confidence: base.confidence * 1.15,
..base
}),
}
}
fn renorm(w: StageWeights) -> StageWeights {
let sum = w.relevance + w.confidence + w.freshness + w.scope;
if sum <= f32::EPSILON {
return w;
}
StageWeights {
relevance: w.relevance / sum,
confidence: w.confidence / sum,
freshness: w.freshness / sum,
scope: w.scope / sum,
}
}
fn task_kind_prefer_roles(kind: TaskKind) -> &'static [&'static str] {
match kind {
TaskKind::Feature => &[],
TaskKind::Debug => &["failure_pattern"],
TaskKind::Refactor => &["convention"],
TaskKind::Investigation => &["fact", "preference"],
TaskKind::Docs => &["convention"],
}
}
use time::OffsetDateTime;
use crate::embeddings::{
self, DEFAULT_HYBRID_ALPHA, Embedder, cosine_similarity, decode_embedding,
};
#[derive(Debug, Clone)]
pub(crate) struct QueryEmbedding {
pub(crate) vector: Vec<f32>,
pub(crate) model_id: String,
}
impl QueryEmbedding {
fn from_embedder(embedder: &dyn Embedder, query: &str) -> Option<Self> {
if embedder.is_noop() {
return None;
}
match embedder.embed(query) {
Ok(v) if v.len() == embedder.dim() => Some(Self {
vector: v,
model_id: embedder.model_id().to_string(),
}),
_ => None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextCapsule {
pub id: String,
pub kind: String,
pub summary: String,
pub token_estimate: u32,
pub expansion_handle: String,
pub provenance: Vec<ProvenanceRef>,
pub confidence: f32,
pub freshness: f32,
pub relevance: f32,
pub scope_weight: f32,
pub score: f32,
#[serde(default)]
pub superseded_hint: bool,
#[serde(default)]
pub rerank_policy_tier: i8,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub claim_revision: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub facts: Vec<crate::fact_store::StoredFact>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub rerank_usefulness: Option<f32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub rerank_trust: Option<f32>,
}
pub fn memory_revision_bindings(
capsules: &[ContextCapsule],
) -> std::collections::BTreeMap<String, String> {
let mut bindings = std::collections::BTreeMap::new();
let mut ambiguous = std::collections::HashSet::new();
for c in capsules {
if let (Some(id), Some(revision)) = (
c.expansion_handle.strip_prefix("memory:"),
c.claim_revision.as_ref(),
) {
if bindings.get(id).is_some_and(|old| old != revision) {
ambiguous.insert(id.to_string());
}
bindings.insert(id.to_string(), revision.clone());
}
}
for id in ambiguous {
bindings.remove(&id);
}
bindings
}
impl ContextCapsule {
pub fn wire_minimal(summary: String, kind: String, score: f32) -> Self {
Self {
id: String::new(),
kind,
summary,
token_estimate: 0,
expansion_handle: String::new(),
provenance: Vec::new(),
confidence: 0.0,
freshness: 0.0,
relevance: 0.0,
scope_weight: 0.0,
score,
superseded_hint: false,
rerank_policy_tier: 0,
claim_revision: None,
facts: vec![],
rerank_usefulness: None,
rerank_trust: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProvenanceRef {
pub source: String,
pub id: String,
pub excerpt: Option<String>,
}
#[derive(Debug, Clone, Default)]
pub struct ContextRequest {
pub include_fact_evidence: bool,
pub defer_fact_budget: bool,
pub stage: String,
pub query: String,
pub budget_tokens: u32,
pub fusion: String,
pub normalization: String,
pub tags: Vec<String>,
pub min_score: f32,
pub max_capsules: usize,
pub prefer_roles: Vec<String>,
pub kinds: Vec<String>,
pub min_semantic_score: f32,
pub min_semantic_score_override: Option<f32>,
pub min_lexical_coverage: f32,
pub min_lexical_coverage_override: Option<f32>,
pub task_kind: TaskKind,
pub abstain_evidence: f32,
pub abstain_evidence_override: Option<f32>,
}
#[derive(Debug, Clone)]
pub struct ContextBundle {
pub stage: String,
pub budget_tokens: u32,
pub used_tokens: u32,
pub capsules: Vec<ContextCapsule>,
pub excluded: Vec<ContextCapsule>,
pub skipped: bool,
pub top_score: f32,
pub top_abs_evidence: f32,
pub evidence_coverage: f32,
pub uncovered_terms: Vec<String>,
pub chronological: bool,
pub known_fact_conflicts: Vec<String>,
}
fn coverage_token_idf(conn: &Connection, tokens: &[String]) -> KimetsuResult<HashMap<String, f32>> {
token_idf(conn, tokens, false)
}
pub fn partial_evidence_notice(bundle: &ContextBundle) -> Option<String> {
if bundle.skipped || bundle.capsules.is_empty() {
return None; }
if bundle.evidence_coverage > PARTIAL_EVIDENCE_COVERAGE || bundle.uncovered_terms.is_empty() {
return None;
}
const MAX_NAMED: usize = 6;
let named: Vec<&str> = bundle
.uncovered_terms
.iter()
.take(MAX_NAMED)
.map(String::as_str)
.collect();
let more = bundle.uncovered_terms.len().saturating_sub(named.len());
let suffix = if more > 0 {
format!(" (and {more} more)")
} else {
String::new()
};
Some(format!(
"Partial memory: nothing above covers {}{}. Treat the rest as unknown \
rather than inferring it.",
named.join(", "),
suffix
))
}
pub const PARTIAL_EVIDENCE_COVERAGE: f32 = 0.5;
pub(crate) fn evidence_coverage(
conn: &Connection,
query: &str,
capsules: &[ContextCapsule],
) -> (f32, Vec<String>) {
let content = content_tokens(query);
if content.is_empty() {
return (1.0, Vec::new());
}
let Ok(idf) = coverage_token_idf(conn, &content) else {
return (1.0, Vec::new());
};
let haystack = capsules
.iter()
.map(|c| c.summary.to_ascii_lowercase())
.collect::<Vec<_>>()
.join(" ");
let mut total = 0.0f32;
let mut hit = 0.0f32;
let mut uncovered = Vec::new();
for token in &content {
let weight = idf.get(token).copied().unwrap_or(0.0);
if weight <= 0.0 {
continue; }
total += weight;
if haystack.contains(token.as_str()) {
hit += weight;
} else {
uncovered.push(token.clone());
}
}
if total <= f32::EPSILON {
return (1.0, Vec::new());
}
(hit / total, uncovered)
}
#[derive(Debug, Clone)]
pub(crate) struct Candidate {
pub(crate) capsule: ContextCapsule,
pub(crate) raw_relevance: f32,
pub(crate) embedding: Option<Vec<f32>>,
pub(crate) cosine: Option<f32>,
pub(crate) created_at: Option<String>,
}
pub fn retrieve_context(
conn: &Connection,
repo_root: &str,
weights: &BrokerWeights,
request: ContextRequest,
) -> KimetsuResult<ContextBundle> {
retrieve_context_multi(conn, repo_root, weights, request, &[])
}
pub fn retrieve_context_multi(
conn: &Connection,
repo_root: &str,
weights: &BrokerWeights,
request: ContextRequest,
extra_memory_conns: &[&Connection],
) -> KimetsuResult<ContextBundle> {
let embedder = embeddings::open_default_embedder();
retrieve_context_with_embedder(
conn,
repo_root,
weights,
request,
extra_memory_conns,
embedder,
)
}
pub fn retrieve_context_with_embedder(
conn: &Connection,
repo_root: &str,
weights: &BrokerWeights,
request: ContextRequest,
extra_memory_conns: &[&Connection],
embedder: &dyn Embedder,
) -> KimetsuResult<ContextBundle> {
retrieve_context_with_embedder_and_backend(
conn,
repo_root,
weights,
request,
extra_memory_conns,
embedder,
&crate::backend::FlatBackend {
fusion: crate::fusion::Fusion::Linear,
},
)
}
pub(crate) fn retrieve_context_with_embedder_and_backend(
conn: &Connection,
repo_root: &str,
weights: &BrokerWeights,
request: ContextRequest,
extra_memory_conns: &[&Connection],
embedder: &dyn Embedder,
backend: &dyn crate::backend::RetrievalBackend,
) -> KimetsuResult<ContextBundle> {
let query_embedding = QueryEmbedding::from_embedder(embedder, &request.query);
let half_life_days = weights.decay_half_life_days;
let mut candidates = Vec::new();
candidates.extend(backend.memory_candidates(
conn,
&request.query,
query_embedding.as_ref(),
half_life_days,
request.include_fact_evidence,
)?);
for extra in extra_memory_conns {
candidates.extend(backend.memory_candidates(
extra,
&request.query,
query_embedding.as_ref(),
half_life_days,
request.include_fact_evidence,
)?);
}
crate::reinforce::apply_query_routing(
conn,
&request.query,
query_embedding.as_ref(),
&mut candidates,
);
candidates.extend(repo_file_candidates(conn, repo_root, &request.query, 30)?);
candidates.extend(manifest_candidates(conn, repo_root, &request.query)?);
if !request.kinds.is_empty() {
candidates.retain(|c| {
request
.kinds
.iter()
.any(|k| capsule_matches_kind(&c.capsule, k))
});
}
if request.min_lexical_coverage > 0.0 {
let content = content_tokens(&request.query);
if !content.is_empty() {
let idf = corpus_token_idf(conn, &content)?;
let total_idf: f32 = content
.iter()
.map(|t| idf.get(t).copied().unwrap_or(0.0))
.sum();
if total_idf > f32::EPSILON {
candidates.retain(|c| {
if c.capsule.kind != "memory" {
return true; }
if c.cosine.is_some_and(|cos| cos >= SEMANTIC_KEEP_COSINE) {
return true;
}
weighted_coverage(&content, &idf, &c.capsule.summary)
>= request.min_lexical_coverage
});
}
}
}
let stage_weights = weights_for_stage(weights, &request.stage);
let effective_weights = weights_for_task_kind(stage_weights, request.task_kind);
normalize_and_score(
&mut candidates,
effective_weights,
Normalization::from_config(&request.normalization),
);
let kind_role_hints = task_kind_prefer_roles(request.task_kind);
let mut effective_prefer_roles: Vec<String> = request.prefer_roles.clone();
for &hint in kind_role_hints {
let hint_s = hint.to_string();
if !effective_prefer_roles.contains(&hint_s) {
effective_prefer_roles.push(hint_s);
}
}
if !request.tags.is_empty() || !effective_prefer_roles.is_empty() {
let tags_lc: Vec<String> = request
.tags
.iter()
.map(|t| t.to_ascii_lowercase())
.collect();
for c in &mut candidates {
let summary_lc = c.capsule.summary.to_ascii_lowercase();
if !tags_lc.is_empty() && tags_lc.iter().any(|t| summary_lc.contains(t.as_str())) {
c.capsule.score *= 1.4;
}
if !effective_prefer_roles.is_empty()
&& effective_prefer_roles.iter().any(|r| {
if c.capsule.kind == "memory" {
capsule_matches_kind(&c.capsule, r.as_str())
} else {
c.capsule.kind.contains(r.as_str())
}
})
{
c.capsule.score *= 1.3;
}
}
}
apply_supersession_penalty(&mut candidates);
if query_embedding.is_some() && request.min_semantic_score > 0.0 {
candidates.retain(|c| {
match c.cosine {
Some(cos) => cos >= request.min_semantic_score,
None => true,
}
});
}
candidates.sort_by(|a, b| {
b.capsule
.score
.partial_cmp(&a.capsule.score)
.unwrap_or(Ordering::Equal)
.then_with(|| {
b.capsule
.freshness
.partial_cmp(&a.capsule.freshness)
.unwrap_or(Ordering::Equal)
})
.then_with(|| a.capsule.expansion_handle.cmp(&b.capsule.expansion_handle))
});
let embedding_mmr_ran = query_embedding.is_some() && !candidates.is_empty();
let candidates = if embedding_mmr_ran {
apply_candidate_mmr_diversity(candidates, 0.7)
} else {
candidates
};
let created_at_by_handle: std::collections::HashMap<String, String> =
if crate::ordering::is_ordering_query(&request.query) {
candidates
.iter()
.filter_map(|c| {
c.created_at
.clone()
.map(|ts| (c.capsule.expansion_handle.clone(), ts))
})
.collect()
} else {
std::collections::HashMap::new()
};
let top_abs_evidence = candidates
.iter()
.filter(|c| c.capsule.kind == "memory")
.filter_map(|c| c.cosine)
.fold(f32::NAN, f32::max);
let top_abs_evidence = if top_abs_evidence.is_nan() {
-1.0
} else {
top_abs_evidence
};
let memory_only = candidates.iter().all(|c| c.capsule.kind == "memory");
let mut capsules = candidates
.into_iter()
.map(|candidate| candidate.capsule)
.collect::<Vec<_>>();
if !embedding_mmr_ran {
capsules.sort_by(|left, right| {
right
.score
.partial_cmp(&left.score)
.unwrap_or(Ordering::Equal)
.then_with(|| {
right
.freshness
.partial_cmp(&left.freshness)
.unwrap_or(Ordering::Equal)
})
.then_with(|| left.expansion_handle.cmp(&right.expansion_handle))
});
}
let top_score = capsules.first().map(|c| c.score).unwrap_or(0.0);
let composite_skip = request.min_score > 0.0 && top_score < request.min_score;
let evidence_skip = request.abstain_evidence > 0.0
&& memory_only
&& top_abs_evidence >= 0.0
&& top_abs_evidence < (request.abstain_evidence - abstain_band_width()).max(0.0);
if composite_skip || evidence_skip {
return Ok(ContextBundle {
stage: request.stage,
budget_tokens: request.budget_tokens,
used_tokens: 0,
capsules: Vec::new(),
excluded: capsules,
skipped: true,
top_score,
top_abs_evidence,
evidence_coverage: 0.0,
uncovered_terms: Vec::new(),
chronological: false,
known_fact_conflicts: vec![],
});
}
let capsules = apply_mmr_diversity(capsules, 0.7);
let capsule_budget = request.budget_tokens / 2;
let mut used_tokens = 0u32;
let mut included = Vec::new();
let mut excluded = Vec::new();
for capsule in capsules {
if request.max_capsules > 0 && included.len() >= request.max_capsules {
excluded.push(capsule);
continue;
}
if (request.defer_fact_budget && request.max_capsules > 0)
|| used_tokens.saturating_add(capsule.token_estimate) <= capsule_budget
{
used_tokens = used_tokens.saturating_add(capsule.token_estimate);
included.push(capsule);
} else {
excluded.push(capsule);
}
}
let (coverage, uncovered_terms) = evidence_coverage(conn, &request.query, &included);
let chronological = !created_at_by_handle.is_empty();
let included = if chronological {
let dated = crate::ordering::render_chronologically(included, &created_at_by_handle);
used_tokens = dated.iter().map(|c| c.token_estimate).sum();
dated
} else {
included
};
Ok(ContextBundle {
stage: request.stage,
budget_tokens: request.budget_tokens,
used_tokens,
capsules: included,
excluded,
skipped: false,
top_score,
top_abs_evidence,
evidence_coverage: coverage,
uncovered_terms,
chronological,
known_fact_conflicts: vec![],
})
}
pub fn search_memories_including_expired(
conn: &Connection,
limit: u32,
) -> KimetsuResult<Vec<ContextCapsule>> {
let mut stmt = conn.prepare_cached(
"
SELECT memory_id, scope, kind, text, confidence, created_at,
use_count, usefulness_score, valid_from, valid_to
FROM memories
WHERE invalidated_at IS NULL
AND superseded_by IS NULL
ORDER BY created_at DESC
LIMIT ?1
",
)?;
let rows = stmt.query_map(params![limit], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, String>(3)?,
row.get::<_, f32>(4)?,
row.get::<_, String>(5)?,
row.get::<_, i64>(6)?,
row.get::<_, f64>(7)?,
row.get::<_, Option<String>>(8)?,
row.get::<_, Option<String>>(9)?,
))
})?;
let now_utc = OffsetDateTime::now_utc();
let mut capsules = Vec::new();
for row in rows {
let (
memory_id,
scope,
kind,
text,
confidence,
created_at,
_use_count,
_usefulness,
_valid_from,
valid_to,
) = row?;
let freshness = freshness(&created_at);
let scope_weight = scope_weight(&scope);
let suffix = if let Some(ref vt) = valid_to {
if OffsetDateTime::parse(vt, &time::format_description::well_known::Rfc3339)
.is_ok_and(|end| end <= now_utc)
{
format!(" [expired valid_to={vt}]")
} else {
format!(" [valid_to={vt}]")
}
} else {
String::new()
};
let revision = crate::projector::claim_revision_at(conn, &memory_id, None)?;
let facts = crate::fact_store::load(conn, &memory_id, &revision)?;
let claim_revision = Some(revision);
capsules.push(ContextCapsule {
id: new_id().to_string(),
kind: "memory".to_string(),
summary: format!("{scope}:{kind} - {text}{suffix}"),
token_estimate: estimate_tokens(&text) + 8,
expansion_handle: format!("memory:{memory_id}"),
provenance: vec![ProvenanceRef {
source: "Memory".to_string(),
id: memory_id,
excerpt: Some(excerpt(&text)),
}],
confidence,
freshness,
relevance: 0.0,
scope_weight,
score: 0.0,
superseded_hint: false,
rerank_policy_tier: 0,
claim_revision,
facts,
rerank_usefulness: None,
rerank_trust: None,
});
}
Ok(capsules)
}
pub fn search_repo_files(
conn: &Connection,
repo_root: &str,
query: &str,
limit: u32,
) -> KimetsuResult<Vec<ContextCapsule>> {
let candidates = repo_file_candidates(conn, repo_root, query, limit)?;
let mut capsules = candidates
.into_iter()
.map(|mut candidate| {
candidate.capsule.relevance = candidate.raw_relevance;
candidate.capsule.score = candidate.raw_relevance;
candidate.capsule
})
.collect::<Vec<_>>();
capsules.sort_by(|left, right| {
right
.score
.partial_cmp(&left.score)
.unwrap_or(Ordering::Equal)
.then_with(|| left.expansion_handle.cmp(&right.expansion_handle))
});
Ok(capsules)
}
#[cfg(feature = "embeddings")]
fn memory_ann_candidates(
conn: &Connection,
qe: &QueryEmbedding,
k: u32,
query_tokens: &[String],
half_life_days: f32,
include_facts: bool,
) -> KimetsuResult<Vec<Candidate>> {
let handle = crate::ann::handle_for_query(conn, qe.vector.len(), &qe.model_id)?;
let hits = handle
.read()
.unwrap_or_else(|p| p.into_inner())
.search(&qe.vector, k as usize)?;
let knn_rowids: Vec<i64> = hits.into_iter().map(|(rowid, _dist)| rowid).collect();
if knn_rowids.is_empty() {
return Ok(Vec::new());
}
let placeholders: String = knn_rowids
.iter()
.enumerate()
.map(|(i, _)| format!("?{}", i + 1))
.collect::<Vec<_>>()
.join(", ");
let sql = format!(
"SELECT memory_id, scope, kind, text, confidence, created_at,
use_count, usefulness_score, embedding, embedding_model,
last_useful_at, provenance_snapshot_json
FROM memories
WHERE invalidated_at IS NULL
AND superseded_by IS NULL
AND (valid_from IS NULL OR julianday(valid_from) <= julianday('now'))
AND (valid_to IS NULL OR julianday(valid_to) > julianday('now'))
AND embedding_model = ?{model_param}
AND rowid IN ({placeholders})",
model_param = knn_rowids.len() + 1
);
let mut stmt = conn.prepare(&sql)?;
let mut params_vec: Vec<&dyn rusqlite::ToSql> = knn_rowids
.iter()
.map(|n| n as &dyn rusqlite::ToSql)
.collect();
params_vec.push(&qe.model_id);
let rows_iter = stmt.query_map(params_vec.as_slice(), |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, String>(3)?,
row.get::<_, f32>(4)?,
row.get::<_, String>(5)?,
row.get::<_, i64>(6)?,
row.get::<_, f64>(7)?,
row.get::<_, Option<Vec<u8>>>(8)?,
row.get::<_, Option<String>>(9)?,
row.get::<_, Option<String>>(10)?,
row.get::<_, Option<String>>(11)?,
))
})?;
let mut candidates = Vec::new();
for row in rows_iter {
let (
memory_id,
scope,
kind,
text,
confidence,
created_at,
use_count,
usefulness_score,
embedding,
embedding_model,
last_useful_at,
provenance_snapshot,
) = row?;
let (cosine, row_vec) =
compute_cosine_and_vec(Some(qe), embedding.as_deref(), embedding_model.as_deref());
let claim_revision = crate::projector::claim_revision_at(conn, &memory_id, None)?;
if let Some(mut candidate) = memory_row_to_candidate(
query_tokens,
memory_id,
scope,
kind,
text,
confidence,
created_at,
use_count,
usefulness_score,
last_useful_at,
provenance_snapshot,
half_life_days,
None, cosine,
row_vec,
) {
candidate.capsule.claim_revision = Some(claim_revision);
hydrate_fact_evidence(conn, &mut candidate, include_facts)?;
candidates.push(candidate);
}
}
Ok(candidates)
}
pub(crate) fn hydrate_fact_evidence(
conn: &Connection,
candidate: &mut Candidate,
include_facts: bool,
) -> KimetsuResult<()> {
if include_facts {
if let (Some(id), Some(revision)) = (
candidate.capsule.expansion_handle.strip_prefix("memory:"),
candidate.capsule.claim_revision.as_deref(),
) {
candidate.capsule.facts = crate::fact_store::load(conn, id, revision)?;
}
}
Ok(())
}
pub(crate) fn memory_candidates_flat(
conn: &Connection,
query: &str,
query_embedding: Option<&QueryEmbedding>,
half_life_days: f32,
fusion: crate::fusion::Fusion,
include_facts: bool,
) -> KimetsuResult<Vec<Candidate>> {
memory_candidates(
conn,
query,
query_embedding,
half_life_days,
fusion,
include_facts,
)
}
fn memory_candidates(
conn: &Connection,
query: &str,
query_embedding: Option<&QueryEmbedding>,
half_life_days: f32,
#[cfg_attr(not(feature = "embeddings"), allow(unused_variables))] fusion: crate::fusion::Fusion,
include_facts: bool,
) -> KimetsuResult<Vec<Candidate>> {
let query_tokens = query_tokens(query);
#[cfg(feature = "embeddings")]
if let Some(qe) = query_embedding {
let fts_candidates = if let Some(fts_query) = fts_query(query) {
memory_fts_candidates(
conn,
&query_tokens,
&fts_query,
80,
Some(qe),
half_life_days,
include_facts,
)?
} else {
Vec::new()
};
let ann_candidates =
memory_ann_candidates(conn, qe, 80, &query_tokens, half_life_days, include_facts)?;
return Ok(crate::fusion::fuse(
fusion,
vec![fts_candidates, ann_candidates],
));
}
if let Some(fts_query) = fts_query(query) {
let candidates = memory_fts_candidates(
conn,
&query_tokens,
&fts_query,
80,
query_embedding,
half_life_days,
include_facts,
)?;
if !candidates.is_empty() {
return Ok(candidates);
}
}
latest_memory_candidates(
conn,
&query_tokens,
200,
query_embedding,
half_life_days,
include_facts,
)
}
fn latest_memory_candidates(
conn: &Connection,
query_tokens: &[String],
limit: u32,
query_embedding: Option<&QueryEmbedding>,
half_life_days: f32,
include_facts: bool,
) -> KimetsuResult<Vec<Candidate>> {
let mut stmt = conn.prepare_cached(
"
SELECT memory_id, scope, kind, text, confidence, created_at,
use_count, usefulness_score, embedding, embedding_model,
last_useful_at, provenance_snapshot_json
FROM memories
WHERE invalidated_at IS NULL
AND superseded_by IS NULL
AND (valid_from IS NULL OR julianday(valid_from) <= julianday('now'))
AND (valid_to IS NULL OR julianday(valid_to) > julianday('now'))
ORDER BY created_at DESC
LIMIT ?1
",
)?;
let rows = stmt.query_map(params![limit], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, String>(3)?,
row.get::<_, f32>(4)?,
row.get::<_, String>(5)?,
row.get::<_, i64>(6)?,
row.get::<_, f64>(7)?,
row.get::<_, Option<Vec<u8>>>(8)?,
row.get::<_, Option<String>>(9)?,
row.get::<_, Option<String>>(10)?,
row.get::<_, Option<String>>(11)?,
))
})?;
let mut candidates = Vec::new();
for row in rows {
let (
memory_id,
scope,
kind,
text,
confidence,
created_at,
use_count,
usefulness_score,
embedding,
embedding_model,
last_useful_at,
provenance_snapshot,
) = row?;
let (cosine, row_vec) = compute_cosine_and_vec(
query_embedding,
embedding.as_deref(),
embedding_model.as_deref(),
);
let claim_revision = crate::projector::claim_revision_at(conn, &memory_id, None)?;
if let Some(mut candidate) = memory_row_to_candidate(
query_tokens,
memory_id,
scope,
kind,
text,
confidence,
created_at,
use_count,
usefulness_score,
last_useful_at,
provenance_snapshot,
half_life_days,
None,
cosine,
row_vec,
) {
candidate.capsule.claim_revision = Some(claim_revision);
hydrate_fact_evidence(conn, &mut candidate, include_facts)?;
candidates.push(candidate);
}
}
Ok(candidates)
}
fn memory_fts_candidates(
conn: &Connection,
query_tokens: &[String],
fts_query: &str,
limit: u32,
query_embedding: Option<&QueryEmbedding>,
half_life_days: f32,
include_facts: bool,
) -> KimetsuResult<Vec<Candidate>> {
let mut stmt = conn.prepare_cached(
"
SELECT m.memory_id, m.scope, m.kind, m.text, m.confidence, m.created_at,
m.use_count, m.usefulness_score, bm25(memories_fts) AS rank,
m.embedding, m.embedding_model, m.last_useful_at,
m.provenance_snapshot_json
FROM memories_fts
JOIN memories m
ON m.memory_id = memories_fts.memory_id
WHERE m.invalidated_at IS NULL
AND m.superseded_by IS NULL
AND (m.valid_from IS NULL OR julianday(m.valid_from) <= julianday('now'))
AND (m.valid_to IS NULL OR julianday(m.valid_to) > julianday('now'))
AND memories_fts MATCH ?1
ORDER BY rank
LIMIT ?2
",
)?;
let rows = stmt.query_map(params![fts_query, limit], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, String>(3)?,
row.get::<_, f32>(4)?,
row.get::<_, String>(5)?,
row.get::<_, i64>(6)?,
row.get::<_, f64>(7)?,
row.get::<_, f64>(8)?,
row.get::<_, Option<Vec<u8>>>(9)?,
row.get::<_, Option<String>>(10)?,
row.get::<_, Option<String>>(11)?,
row.get::<_, Option<String>>(12)?,
))
})?;
let mut candidates = Vec::new();
for row in rows {
let (
memory_id,
scope,
kind,
text,
confidence,
created_at,
use_count,
usefulness_score,
rank,
embedding,
embedding_model,
last_useful_at,
provenance_snapshot,
) = row?;
let fts_relevance = (-rank as f32).max(0.0);
let (cosine, row_vec) = compute_cosine_and_vec(
query_embedding,
embedding.as_deref(),
embedding_model.as_deref(),
);
let claim_revision = crate::projector::claim_revision_at(conn, &memory_id, None)?;
if let Some(mut candidate) = memory_row_to_candidate(
query_tokens,
memory_id,
scope,
kind,
text,
confidence,
created_at,
use_count,
usefulness_score,
last_useful_at,
provenance_snapshot,
half_life_days,
Some(fts_relevance),
cosine,
row_vec,
) {
candidate.capsule.claim_revision = Some(claim_revision);
hydrate_fact_evidence(conn, &mut candidate, include_facts)?;
candidates.push(candidate);
}
}
Ok(candidates)
}
fn compute_cosine_and_vec(
query_embedding: Option<&QueryEmbedding>,
row_bytes: Option<&[u8]>,
row_model: Option<&str>,
) -> (Option<f32>, Option<Vec<f32>>) {
let q = match query_embedding {
Some(q) => q,
None => return (None, None),
};
let bytes = match row_bytes {
Some(b) => b,
None => return (None, None),
};
let model = match row_model {
Some(m) => m,
None => return (None, None),
};
if model != q.model_id {
return (None, None);
}
let row_vec = match decode_embedding(bytes, Some(q.vector.len())) {
Ok(v) => v,
Err(_) => return (None, None),
};
let score = cosine_similarity(&q.vector, &row_vec);
(Some(score), Some(row_vec))
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn memory_row_to_candidate(
query_tokens: &[String],
memory_id: String,
scope: String,
kind: String,
text: String,
confidence: f32,
created_at: String,
use_count: i64,
usefulness_score: f64,
last_useful_at: Option<String>,
provenance_snapshot: Option<String>,
half_life_days: f32,
raw_relevance_override: Option<f32>,
cosine_score: Option<f32>,
row_embedding: Option<Vec<f32>>,
) -> Option<Candidate> {
let lexical = lexical_relevance(query_tokens, &format!("{kind} {text}"));
let lexical_term = raw_relevance_override.unwrap_or(lexical).max(lexical);
let raw_relevance = match cosine_score {
Some(c) => {
let normalized_cos = ((c + 1.0) * 0.5).clamp(0.0, 1.0);
(1.0 - DEFAULT_HYBRID_ALPHA) * lexical_term + DEFAULT_HYBRID_ALPHA * normalized_cos
}
None => lexical_term,
};
if raw_relevance <= 0.0 && !query_tokens.is_empty() {
return None;
}
let freshness = freshness(&created_at);
let scope_weight = scope_weight(&scope);
let raw_multiplier = usefulness_multiplier(usefulness_score as f32, use_count as u32);
let decay = usefulness_decay(last_useful_at.as_deref(), &created_at, half_life_days);
let multiplier = 1.0 + (raw_multiplier - 1.0) * decay;
let biased_relevance = apply_usefulness_boost(raw_relevance, multiplier);
let rerank_policy_tier = if multiplier > 1.0 + f32::EPSILON {
1
} else if multiplier < 1.0 - f32::EPSILON {
-1
} else {
0
};
let provenance =
crate::trust::Provenance::from_snapshot(provenance_snapshot.as_deref().unwrap_or("{}"));
let trusted_relevance =
biased_relevance * crate::trust::trust_multiplier(provenance, last_useful_at.is_some());
Some(Candidate {
raw_relevance: trusted_relevance,
embedding: row_embedding,
cosine: cosine_score,
created_at: Some(created_at),
capsule: ContextCapsule {
id: new_id().to_string(),
kind: "memory".to_string(),
summary: format!("{scope}:{kind} - {text}"),
token_estimate: estimate_tokens(&text) + 8,
expansion_handle: format!("memory:{memory_id}"),
provenance: vec![ProvenanceRef {
source: "Memory".to_string(),
id: memory_id,
excerpt: Some(excerpt(&text)),
}],
confidence,
freshness,
relevance: 0.0,
scope_weight,
score: 0.0,
superseded_hint: false,
rerank_policy_tier,
claim_revision: None,
facts: vec![],
rerank_usefulness: Some(multiplier),
rerank_trust: Some(crate::trust::trust_multiplier(
provenance,
last_useful_at.is_some(),
)),
},
})
}
pub(crate) fn usefulness_decay(
last_useful_at: Option<&str>,
created_at: &str,
half_life_days: f32,
) -> f32 {
if half_life_days <= 0.0 {
return 1.0;
}
let reference = last_useful_at.unwrap_or(created_at);
let Ok(reference_ts) =
OffsetDateTime::parse(reference, &time::format_description::well_known::Rfc3339)
else {
return 1.0;
};
let age = OffsetDateTime::now_utc() - reference_ts;
let age_days = (age.whole_seconds().max(0) as f32) / 86_400.0;
let exponent = -std::f32::consts::LN_2 * age_days / half_life_days;
exponent.exp().clamp(0.0, 1.0)
}
pub(crate) use crate::scoring::USEFULNESS_BOOST_CAP;
pub(crate) fn apply_usefulness_boost(raw_relevance: f32, multiplier: f32) -> f32 {
if multiplier <= 1.0 {
return raw_relevance * multiplier;
}
(raw_relevance * multiplier).min(raw_relevance + USEFULNESS_BOOST_CAP)
}
pub(crate) fn usefulness_multiplier(usefulness_score: f32, use_count: u32) -> f32 {
use crate::scoring::{FULL_CONFIDENCE_USES, MULTIPLIER_MAX, MULTIPLIER_MIN};
if use_count == 0 {
return 1.0;
}
let ratio = usefulness_score / use_count as f32; let normalized = ((ratio + 1.0) / 2.0).clamp(0.0, 1.0); let full_multiplier = MULTIPLIER_MIN + normalized * (MULTIPLIER_MAX - MULTIPLIER_MIN);
let confidence = (use_count as f32 / FULL_CONFIDENCE_USES as f32).min(1.0);
1.0 * (1.0 - confidence) + full_multiplier * confidence
}
fn repo_file_candidates(
conn: &Connection,
repo_root: &str,
query: &str,
limit: u32,
) -> KimetsuResult<Vec<Candidate>> {
let Some(fts_query) = fts_query(query) else {
return Ok(Vec::new());
};
let mut stmt = conn.prepare_cached(
"
SELECT path, snippet, language_guess, bm25(repo_files_fts) AS rank
FROM repo_files_fts
WHERE repo_root = ?1 AND repo_files_fts MATCH ?2
ORDER BY rank
LIMIT ?3
",
)?;
let rows = stmt.query_map(params![repo_root, fts_query, limit], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, f64>(3)?,
))
})?;
let mut candidates = Vec::new();
for row in rows {
let (path, snippet, language, rank) = row?;
let raw_relevance = (-rank as f32).max(0.0);
let summary = format!("{path} ({language}) - {}", excerpt(&snippet));
let token_estimate = estimate_tokens(&summary) + 8;
candidates.push(Candidate {
raw_relevance,
embedding: None,
cosine: None,
created_at: None,
capsule: ContextCapsule {
id: new_id().to_string(),
kind: "repo_file".to_string(),
summary,
token_estimate,
expansion_handle: format!("file:{path}"),
provenance: vec![ProvenanceRef {
source: "RepoFile".to_string(),
id: path.clone(),
excerpt: Some(excerpt(&snippet)),
}],
confidence: 0.9,
freshness: 1.0,
relevance: 0.0,
scope_weight: 0.9,
score: 0.0,
superseded_hint: false,
rerank_policy_tier: 0,
claim_revision: None,
facts: vec![],
rerank_usefulness: None,
rerank_trust: None,
},
});
}
Ok(candidates)
}
fn manifest_candidates(
conn: &Connection,
repo_root: &str,
query: &str,
) -> KimetsuResult<Vec<Candidate>> {
if let Some(fts_query) = fts_query(query) {
let candidates = manifest_fts_candidates(conn, repo_root, &fts_query, 30)?;
if !candidates.is_empty() {
return Ok(candidates);
}
}
let query_tokens = query_tokens(query);
let mut stmt = conn.prepare_cached(
"
SELECT manifest_path, manifest_kind, parsed_summary_json
FROM repo_manifests
WHERE repo_root = ?1
ORDER BY manifest_path
",
)?;
let rows = stmt.query_map(params![repo_root], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
))
})?;
let mut candidates = Vec::new();
for row in rows {
let (path, kind, summary_json) = row?;
let raw_relevance =
lexical_relevance(&query_tokens, &format!("{path} {kind} {summary_json}"));
if raw_relevance <= 0.0 && !query_tokens.is_empty() {
continue;
}
let summary = format!("{path} manifest ({kind})");
let token_estimate = estimate_tokens(&summary) + 8;
candidates.push(Candidate {
raw_relevance,
embedding: None,
cosine: None,
created_at: None,
capsule: ContextCapsule {
id: new_id().to_string(),
kind: "repo_manifest".to_string(),
summary,
token_estimate,
expansion_handle: format!("file:{path}"),
provenance: vec![ProvenanceRef {
source: "Manifest".to_string(),
id: path,
excerpt: Some(excerpt(&summary_json)),
}],
confidence: 0.95,
freshness: 1.0,
relevance: 0.0,
scope_weight: 0.9,
score: 0.0,
superseded_hint: false,
rerank_policy_tier: 0,
claim_revision: None,
facts: vec![],
rerank_usefulness: None,
rerank_trust: None,
},
});
}
Ok(candidates)
}
fn manifest_fts_candidates(
conn: &Connection,
repo_root: &str,
fts_query: &str,
limit: u32,
) -> KimetsuResult<Vec<Candidate>> {
let mut stmt = conn.prepare_cached(
"
SELECT manifest_path, manifest_kind, parsed_summary_json,
bm25(repo_manifests_fts) AS rank
FROM repo_manifests_fts
WHERE repo_root = ?1 AND repo_manifests_fts MATCH ?2
ORDER BY rank
LIMIT ?3
",
)?;
let rows = stmt.query_map(params![repo_root, fts_query, limit], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, f64>(3)?,
))
})?;
let mut candidates = Vec::new();
for row in rows {
let (path, kind, summary_json, rank) = row?;
let raw_relevance = (-rank as f32).max(0.0);
let summary = format!("{path} manifest ({kind})");
let token_estimate = estimate_tokens(&summary) + 8;
candidates.push(Candidate {
raw_relevance,
embedding: None,
cosine: None,
created_at: None,
capsule: ContextCapsule {
id: new_id().to_string(),
kind: "repo_manifest".to_string(),
summary,
token_estimate,
expansion_handle: format!("file:{path}"),
provenance: vec![ProvenanceRef {
source: "Manifest".to_string(),
id: path,
excerpt: Some(excerpt(&summary_json)),
}],
confidence: 0.95,
freshness: 1.0,
relevance: 0.0,
scope_weight: 0.9,
score: 0.0,
superseded_hint: false,
rerank_policy_tier: 0,
claim_revision: None,
facts: vec![],
rerank_usefulness: None,
rerank_trust: None,
},
});
}
Ok(candidates)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Normalization {
#[default]
PerKind,
Global,
}
impl Normalization {
pub fn from_config(value: &str) -> Self {
match value.trim().to_ascii_lowercase().as_str() {
"global" => Self::Global,
_ => Self::PerKind,
}
}
}
fn normalize_and_score(
candidates: &mut [Candidate],
weights: StageWeights,
normalization: Normalization,
) {
let mut max_by_kind = HashMap::<String, f32>::new();
let bucket = |candidate: &Candidate| match normalization {
Normalization::PerKind => candidate.capsule.kind.clone(),
Normalization::Global => String::new(),
};
for candidate in candidates.iter() {
max_by_kind
.entry(bucket(candidate))
.and_modify(|max| *max = (*max).max(candidate.raw_relevance))
.or_insert(candidate.raw_relevance);
}
for candidate in candidates {
let max = max_by_kind.get(&bucket(candidate)).copied().unwrap_or(0.0);
let relevance = if max <= f32::EPSILON {
if candidate.raw_relevance > 0.0 {
1.0
} else {
0.0
}
} else {
(candidate.raw_relevance / max).clamp(0.0, 1.0)
};
candidate.capsule.relevance = relevance;
candidate.capsule.score = weights.relevance * relevance
+ weights.confidence * candidate.capsule.confidence
+ weights.freshness * candidate.capsule.freshness
+ weights.scope * candidate.capsule.scope_weight;
}
}
pub(crate) const SUPERSESSION_MIN_COSINE: f32 = 0.85;
pub(crate) const SUPERSESSION_CUE_MIN_COSINE: f32 = 0.82;
pub(crate) const SUPERSESSION_CUE_MIN_SHARED_TOKENS: usize = 3;
pub(crate) const SUPERSESSION_CUE_MIN_LEXICAL_OVERLAP: f32 = 0.35;
pub(crate) const SUPERSESSION_PENALTY: f32 = 0.80;
pub(crate) const SUPERSESSION_MIN_AGE_GAP_SECS: f64 = 1.0;
fn has_explicit_supersession_signal(text: &str) -> bool {
let words = text
.to_ascii_lowercase()
.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { ' ' })
.collect::<String>();
let text = format!(
" {} ",
words.split_whitespace().collect::<Vec<_>>().join(" ")
);
[
" now ",
" no longer ",
" renamed from ",
" switched from ",
" as of ",
" currently ",
" after the ",
" supersedes ",
" replaced by ",
" silently breaks ",
]
.iter()
.any(|signal| text.contains(signal))
}
fn has_supersession_lexical_identity(a: &str, b: &str) -> bool {
let a = content_tokens(a);
let b = content_tokens(b);
let smaller = a.len().min(b.len());
if smaller == 0 {
return false;
}
let shared = a.iter().filter(|token| b.contains(token)).count();
shared >= SUPERSESSION_CUE_MIN_SHARED_TOKENS
&& shared as f32 / smaller as f32 >= SUPERSESSION_CUE_MIN_LEXICAL_OVERLAP
}
pub(crate) fn apply_supersession_penalty(candidates: &mut [Candidate]) {
let parsed: Vec<Option<OffsetDateTime>> = candidates
.iter()
.map(|c| {
if c.capsule.kind != "memory" || c.embedding.is_none() {
return None;
}
c.created_at.as_deref().and_then(|ts| {
OffsetDateTime::parse(ts, &time::format_description::well_known::Rfc3339).ok()
})
})
.collect();
let mut penalized = vec![false; candidates.len()];
for i in 0..candidates.len() {
let Some(ti) = parsed[i] else { continue };
for j in (i + 1)..candidates.len() {
let Some(tj) = parsed[j] else { continue };
let (Some(ei), Some(ej)) = (&candidates[i].embedding, &candidates[j].embedding) else {
continue;
};
let signals = (
has_explicit_supersession_signal(&candidates[i].capsule.summary),
has_explicit_supersession_signal(&candidates[j].capsule.summary),
);
let exactly_one_signal = signals.0 ^ signals.1;
let similarity = crate::embeddings::cosine_similarity(ei, ej);
let cue_related = exactly_one_signal
&& (similarity >= SUPERSESSION_CUE_MIN_COSINE
|| has_supersession_lexical_identity(
&candidates[i].capsule.summary,
&candidates[j].capsule.summary,
));
if similarity < SUPERSESSION_MIN_COSINE && !cue_related {
continue;
}
let gap = (ti - tj).abs();
let subsecond = (gap.whole_milliseconds().unsigned_abs() as f64)
< SUPERSESSION_MIN_AGE_GAP_SECS * 1000.0;
let older = if subsecond {
match signals {
(true, false) => j,
(false, true) => i,
_ => continue,
}
} else {
match ti.cmp(&tj) {
Ordering::Less => i,
Ordering::Greater => j,
Ordering::Equal => continue,
}
};
if similarity < SUPERSESSION_CUE_MIN_COSINE {
let replacement = if older == i { j } else { i };
let replacement_has_signal = if replacement == i {
signals.0
} else {
signals.1
};
let query_favors_replacement =
match (candidates[replacement].cosine, candidates[older].cosine) {
(Some(replacement_cosine), Some(older_cosine)) => {
replacement_cosine >= older_cosine
}
_ => false,
};
if !replacement_has_signal || !query_favors_replacement {
continue;
}
}
if !penalized[older] {
penalized[older] = true;
candidates[older].capsule.score *= SUPERSESSION_PENALTY;
candidates[older].capsule.superseded_hint = true;
}
}
}
}
fn weights_for_stage(weights: &BrokerWeights, stage: &str) -> StageWeights {
match stage {
"localization" => weights.localization.clone(),
"patch_plan" => weights.patch_plan.clone(),
"verification" => weights.verification.clone(),
"review" => weights.review.clone(),
_ => None,
}
.unwrap_or(StageWeights {
relevance: weights.relevance,
confidence: weights.confidence,
freshness: weights.freshness,
scope: weights.scope,
})
}
fn scope_weight(scope: &str) -> f32 {
match scope.parse::<MemoryScope>() {
Ok(MemoryScope::Run) => 1.0,
Ok(MemoryScope::Repo) => 0.9,
Ok(MemoryScope::Project) => 0.7,
Ok(MemoryScope::GlobalUser) => 0.5,
Err(_) => 0.3,
}
}
fn freshness(created_at: &str) -> f32 {
let Ok(created_at) =
OffsetDateTime::parse(created_at, &time::format_description::well_known::Rfc3339)
else {
return 0.5;
};
let age = OffsetDateTime::now_utc() - created_at;
let age_days = age.whole_seconds().max(0) as f32 / 86_400.0;
(-std::f32::consts::LN_2 * age_days / 30.0)
.exp()
.clamp(0.0, 1.0)
}
const SEMANTIC_KEEP_COSINE: f32 = 0.20;
const STOPWORDS: &[&str] = &[
"the", "and", "for", "are", "but", "not", "you", "your", "with", "this", "that", "these",
"those", "from", "into", "about", "what", "whats", "which", "who", "whom", "how", "why",
"when", "where", "can", "could", "would", "should", "will", "shall", "does", "did", "was",
"were", "been", "being", "have", "has", "had", "its", "it", "is", "as", "at", "by", "of", "to",
"in", "on", "or", "an", "be", "do", "me", "my", "we", "us", "our", "im", "ive", "let", "lets",
"please", "tell", "give", "show", "want", "need", "get", "got", "use", "using", "there",
"their", "they", "them", "then", "than", "some", "any", "all", "more", "most", "such", "via",
"per",
"during", "while", "until", "unless", "before", "after", "again", "against", "above", "below",
"between", "through", "under", "over", "because", "also", "just", "only", "very", "much",
"many", "each", "both", "same", "other", "another", "always", "never", "still", "even", "ever",
"every", "first", "found", "thing", "things", "value", "default", "if", "so", "up", "out",
"off", "down", "no", "yes",
];
fn content_tokens(query: &str) -> Vec<String> {
let mut seen = std::collections::HashSet::new();
query
.split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
.map(str::trim)
.filter(|part| part.len() >= 2)
.map(str::to_ascii_lowercase)
.filter(|t| !STOPWORDS.contains(&t.as_str()))
.map(|t| light_stem(&t).to_string())
.filter(|t| seen.insert(t.clone()))
.collect()
}
fn corpus_token_idf(conn: &Connection, tokens: &[String]) -> KimetsuResult<HashMap<String, f32>> {
token_idf(conn, tokens, true)
}
fn token_idf(
conn: &Connection,
tokens: &[String],
zero_absent: bool,
) -> KimetsuResult<HashMap<String, f32>> {
let n: i64 = conn.query_row(
"SELECT COUNT(*) FROM memories_fts JOIN memories m USING(memory_id) WHERE m.invalidated_at IS NULL",
[], |r| r.get(0))?;
let mut idf = HashMap::new();
if n == 0 {
return Ok(idf);
}
let mut stmt = conn.prepare_cached(
"SELECT COUNT(DISTINCT m.memory_id) FROM memories_fts JOIN memories m USING(memory_id)
WHERE memories_fts MATCH ?1 AND m.invalidated_at IS NULL",
)?;
for token in tokens {
if idf.contains_key(token) {
continue;
}
let pattern = format!("text : \"{}\"*", token.replace('"', "\"\""));
let df: i64 = stmt.query_row(params![pattern], |r| r.get(0))?;
let weight = if zero_absent && df == 0 {
0.0
} else {
(((n + 1) as f32) / ((df + 1) as f32)).ln().max(0.0)
};
idf.insert(token.clone(), weight);
}
Ok(idf)
}
fn weighted_coverage(content: &[String], idf: &HashMap<String, f32>, summary: &str) -> f32 {
let haystack = summary.to_ascii_lowercase();
let mut total = 0.0f32;
let mut hit = 0.0f32;
for token in content {
let weight = idf.get(token).copied().unwrap_or(0.0);
total += weight;
if weight > 0.0 && haystack.contains(token.as_str()) {
hit += weight;
}
}
if total <= f32::EPSILON {
0.0
} else {
(hit / total).clamp(0.0, 1.0)
}
}
fn light_stem(token: &str) -> &str {
let mut stem = token;
for suffix in ["ing", "ed", "es", "s"] {
if let Some(stripped) = token.strip_suffix(suffix)
&& stripped.len() >= 4
{
stem = stripped;
break;
}
}
if stem.len() >= 5
&& let Some(trimmed) = stem.strip_suffix('y').or_else(|| stem.strip_suffix('i'))
&& trimmed
.chars()
.next_back()
.is_some_and(|c| !matches!(c, 'a' | 'e' | 'i' | 'o' | 'u'))
{
return trimmed;
}
stem
}
fn query_tokens(query: &str) -> Vec<String> {
let mut tokens: Vec<String> = query
.split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
.map(str::trim)
.filter(|part| part.len() >= 2)
.map(str::to_ascii_lowercase)
.map(|t| light_stem(&t).to_string())
.collect();
let lower = query.to_ascii_lowercase();
for (triggers, expansions) in CLASS_HINTS.iter() {
if triggers.iter().any(|t| lower.contains(t)) {
tokens.extend(expansions.iter().map(|e| e.to_string()));
}
}
tokens
}
const CLASS_HINTS: &[(&[&str], &[&str])] = &[
(
&[
"build",
"compile",
"make",
"cargo",
"cmake",
"configure",
"install",
"train",
"benchmark",
"test suite",
"ray trace",
"render",
],
&[
"shell_background",
"shell_status",
"shell_output",
"shell_stop",
"long_running",
],
),
(
&[
"edit", "modify", "change", "fix", "update", "patch", "refactor", "rename",
],
&["edit_file", "apply_patch", "old_string", "new_string"],
),
(
&[
"read", "inspect", "review", "analyze", "examine", "view", "show",
],
&["read_file", "offset", "limit", "multi_read"],
),
(
&["find", "locate", "search", "look up", "discover", "list"],
&["glob", "search_files", "list_files"],
),
(
&["plan", "step", "checklist", "todo", "task list", "phase"],
&["plan", "todos"],
),
(
&[
"verify",
"check",
"ensure",
"validate",
"pass test",
"verifier",
],
&["finish", "verifier", "verification"],
),
(
&[
"image",
"png",
"jpeg",
"jpg",
"pdf",
"diagram",
"screenshot",
],
&["view_image", "base64", "sha256"],
),
(&["delete", "remove", "rm "], &["delete_file", "recursive"]),
(&["rename", "move file", "mv "], &["move_file"]),
];
fn capsule_matches_kind(capsule: &ContextCapsule, wanted: &str) -> bool {
if capsule.kind == wanted {
return true;
}
if capsule.kind == "memory"
&& let Some((prefix, _)) = capsule.summary.split_once(" - ")
&& let Some((_scope, mkind)) = prefix.split_once(':')
{
return mkind == wanted;
}
false
}
pub(crate) fn fts_query(query: &str) -> Option<String> {
let tokens = query_tokens(query);
if tokens.is_empty() {
return None;
}
Some(
tokens
.into_iter()
.take(12)
.map(|token| format!("{token}*"))
.collect::<Vec<_>>()
.join(" OR "),
)
}
fn apply_candidate_mmr_diversity(mut sorted: Vec<Candidate>, lambda: f32) -> Vec<Candidate> {
if sorted.len() <= 1 {
return sorted;
}
let summaries: Vec<std::collections::HashSet<String>> = sorted
.iter()
.map(|c| summary_token_set(&c.capsule.summary))
.collect();
let mut picked_indices: Vec<usize> = Vec::with_capacity(sorted.len());
let mut remaining: Vec<usize> = (0..sorted.len()).collect();
picked_indices.push(remaining.remove(0));
while !remaining.is_empty() {
let mut best_idx_in_remaining = 0;
let mut best_score = f32::MIN;
for (i, &cand) in remaining.iter().enumerate() {
let mut max_overlap = 0.0f32;
for &p in &picked_indices {
let same_kind = sorted[cand].capsule.kind == sorted[p].capsule.kind;
let raw_overlap = candidate_pair_overlap(
&sorted[cand],
&sorted[p],
&summaries[cand],
&summaries[p],
);
let overlap = if same_kind {
raw_overlap
} else {
raw_overlap * 0.5
};
if overlap > max_overlap {
max_overlap = overlap;
}
}
let mmr = lambda * sorted[cand].capsule.score - (1.0 - lambda) * max_overlap;
if mmr > best_score {
best_score = mmr;
best_idx_in_remaining = i;
}
}
picked_indices.push(remaining.remove(best_idx_in_remaining));
}
let mut taken: Vec<Option<Candidate>> = sorted.drain(..).map(Some).collect();
let mut out = Vec::with_capacity(taken.len());
for idx in picked_indices {
if let Some(c) = taken[idx].take() {
out.push(c);
}
}
out
}
fn candidate_pair_overlap(
a: &Candidate,
b: &Candidate,
tokens_a: &std::collections::HashSet<String>,
tokens_b: &std::collections::HashSet<String>,
) -> f32 {
if let (Some(va), Some(vb)) = (a.embedding.as_deref(), b.embedding.as_deref()) {
cosine_similarity(va, vb).max(0.0)
} else {
jaccard(tokens_a, tokens_b)
}
}
fn apply_mmr_diversity(mut sorted: Vec<ContextCapsule>, lambda: f32) -> Vec<ContextCapsule> {
if sorted.len() <= 1 {
return sorted;
}
let summaries: Vec<std::collections::HashSet<String>> = sorted
.iter()
.map(|c| summary_token_set(&c.summary))
.collect();
let mut picked_indices: Vec<usize> = Vec::with_capacity(sorted.len());
let mut remaining: Vec<usize> = (0..sorted.len()).collect();
picked_indices.push(remaining.remove(0));
while !remaining.is_empty() {
let mut best_idx_in_remaining = 0;
let mut best_score = f32::MIN;
for (i, &cand) in remaining.iter().enumerate() {
let mut max_overlap = 0.0f32;
for &p in &picked_indices {
let raw = jaccard(&summaries[cand], &summaries[p]);
let overlap = if sorted[cand].kind == sorted[p].kind {
raw
} else {
raw * 0.5
};
if overlap > max_overlap {
max_overlap = overlap;
}
}
let mmr = lambda * sorted[cand].score - (1.0 - lambda) * max_overlap;
if mmr > best_score {
best_score = mmr;
best_idx_in_remaining = i;
}
}
picked_indices.push(remaining.remove(best_idx_in_remaining));
}
let mut out = Vec::with_capacity(sorted.len());
let mut taken: Vec<Option<ContextCapsule>> = sorted.drain(..).map(Some).collect();
for idx in picked_indices {
if let Some(c) = taken[idx].take() {
out.push(c);
}
}
out
}
fn summary_token_set(s: &str) -> std::collections::HashSet<String> {
s.split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
.filter(|t| t.len() >= 3)
.map(str::to_ascii_lowercase)
.collect()
}
fn jaccard(a: &std::collections::HashSet<String>, b: &std::collections::HashSet<String>) -> f32 {
if a.is_empty() && b.is_empty() {
return 0.0;
}
let intersection = a.intersection(b).count();
let union = a.union(b).count();
intersection as f32 / union.max(1) as f32
}
fn lexical_relevance(tokens: &[String], haystack: &str) -> f32 {
if tokens.is_empty() {
return 0.0;
}
let haystack = haystack.to_ascii_lowercase();
let matches = tokens
.iter()
.filter(|token| haystack.contains(token.as_str()))
.count();
matches as f32 / tokens.len() as f32
}
pub fn estimate_tokens(text: &str) -> u32 {
((text.split_whitespace().count() as f32) * 1.33).ceil() as u32
}
pub fn compress_for_render(summary: &str, max_sentences: usize) -> String {
if max_sentences == 0 {
return summary.to_string();
}
let text = if let Some(rest) = summary.strip_prefix('[') {
if let Some(idx) = rest.find(']') {
rest[idx + 1..].trim_start()
} else {
summary
}
} else {
summary
};
let text = if let Some(idx) = text.rfind('(') {
let candidate = text[..idx].trim_end();
let inner = &text[idx + 1..];
if inner.contains(':') && inner.trim_end().ends_with(')') {
candidate
} else {
text
}
} else {
text
};
let (scope_prefix, body) = if let Some(dash_pos) = text.find(" - ") {
let prefix_candidate = &text[..dash_pos];
if !prefix_candidate.contains(' ') && prefix_candidate.contains(':') {
let body_start = dash_pos + 3; (&text[..body_start], &text[body_start..])
} else {
("", text)
}
} else {
("", text)
};
let compressed_body = cap_sentences(body, max_sentences);
let result = if scope_prefix.is_empty() {
compressed_body.to_string()
} else {
format!("{scope_prefix}{compressed_body}")
};
if result.trim().is_empty() {
summary.to_string()
} else {
result
}
}
fn cap_sentences(text: &str, n: usize) -> &str {
let bytes = text.as_bytes();
let len = bytes.len();
let mut count = 0;
let mut i = 0;
while i < len {
if bytes[i] == b'.' {
let next = i + 1;
if next < len && (bytes[next] == b' ' || bytes[next] == b'\n') {
count += 1;
if count >= n {
return text[..=i].trim_end();
}
}
}
i += 1;
}
text.trim_end()
}
fn excerpt(text: &str) -> String {
let value = one_line(text);
value.chars().take(256).collect()
}
fn one_line(text: &str) -> String {
text.split_whitespace().collect::<Vec<_>>().join(" ")
}
const FILE_EXPAND_CAP_BYTES: usize = 2048;
pub fn resolve_capsule(
conn: &Connection,
repo_root: &std::path::Path,
handle: &str,
) -> kimetsu_core::KimetsuResult<String> {
if let Some(memory_id) = handle.strip_prefix("memory:") {
let mut stmt = conn.prepare_cached(
"SELECT text FROM memories WHERE memory_id = ? AND invalidated_at IS NULL",
)?;
let text: Option<String> = stmt
.query_row(rusqlite::params![memory_id], |row| row.get(0))
.optional()?;
match text {
Some(t) => Ok(t),
None => {
Err(format!("expand_capsule: no active memory found for handle `{handle}`").into())
}
}
} else if let Some(rel_path) = handle.strip_prefix("file:") {
let path = std::path::Path::new(rel_path);
if path.is_absolute() {
return Err(format!(
"expand_capsule: `{handle}` is an absolute path — only repo-relative paths are supported"
)
.into());
}
for component in path.components() {
match component {
std::path::Component::ParentDir => {
return Err(format!(
"expand_capsule: `{handle}` contains `..` traversal — rejected"
)
.into());
}
std::path::Component::RootDir | std::path::Component::Prefix(_) => {
return Err(format!(
"expand_capsule: `{handle}` is an absolute path — only repo-relative paths are supported"
)
.into());
}
_ => {}
}
}
let full_path = repo_root.join(path);
let bytes = std::fs::read(&full_path)
.map_err(|e| format!("expand_capsule: could not read `{rel_path}`: {e}"))?;
let bounded = if bytes.len() > FILE_EXPAND_CAP_BYTES {
let mut end = FILE_EXPAND_CAP_BYTES;
while end > 0 && (bytes[end] & 0xC0) == 0x80 {
end -= 1;
}
let s = String::from_utf8_lossy(&bytes[..end]);
format!(
"{s}\n[... truncated at {FILE_EXPAND_CAP_BYTES} bytes; call expand_capsule again with a line range if needed]"
)
} else {
String::from_utf8_lossy(&bytes).into_owned()
};
Ok(bounded)
} else if handle.starts_with("run:") {
Err(format!(
"expand_capsule: `run:` handle expansion is not yet supported (handle: `{handle}`)"
)
.into())
} else {
Err(format!(
"expand_capsule: unrecognised handle format `{handle}`; \
expected `memory:<id>`, `file:<path>`, or `run:<id>`"
)
.into())
}
}
pub const ABSTAIN_BAND_WIDTH: f32 = 0.10;
fn abstain_band_width() -> f32 {
std::env::var("KIMETSU_ABSTAIN_BAND_WIDTH")
.ok()
.and_then(|v| v.parse::<f32>().ok())
.unwrap_or(ABSTAIN_BAND_WIDTH)
.clamp(0.0, 1.0)
}
pub const ABSTAIN_RERANK_FLOOR: f32 = 0.9;
fn abstain_rerank_floor() -> f32 {
std::env::var("KIMETSU_ABSTAIN_RERANK_FLOOR")
.ok()
.and_then(|v| v.parse::<f32>().ok())
.unwrap_or(ABSTAIN_RERANK_FLOOR)
}
pub fn rerank_and_arbitrate(
query: &str,
mut bundle: ContextBundle,
reranker: Option<&dyn crate::embeddings::Reranker>,
abstain_evidence: f32,
rerank_floor: f32,
rerank_cap: usize,
) -> ContextBundle {
if bundle.skipped || bundle.capsules.is_empty() {
return bundle;
}
let memory_only = bundle.capsules.iter().all(|c| c.kind == "memory");
let in_band = abstain_evidence > 0.0
&& memory_only
&& bundle.top_abs_evidence >= 0.0
&& bundle.top_abs_evidence < abstain_evidence;
let to_skipped = |mut bundle: ContextBundle| -> ContextBundle {
let rejected = std::mem::take(&mut bundle.capsules);
bundle.excluded.extend(rejected);
bundle.skipped = true;
bundle.used_tokens = 0;
bundle.evidence_coverage = 0.0;
bundle.uncovered_terms = Vec::new();
bundle.chronological = false;
bundle
};
match reranker {
Some(rr) => {
let reranked = rerank_capsules_with_diagnostics(
query,
std::mem::take(&mut bundle.capsules),
rr,
rerank_floor,
rerank_cap,
);
let best_raw_rerank = reranked.best_raw_score.unwrap_or(0.0);
bundle.capsules = reranked.capsules;
bundle.used_tokens = bundle.capsules.iter().map(|c| c.token_estimate).sum();
if in_band && best_raw_rerank < abstain_rerank_floor() {
to_skipped(bundle)
} else {
bundle
}
}
None if in_band => to_skipped(bundle),
None => bundle,
}
}
pub fn rerank_capsules(
query: &str,
capsules: Vec<ContextCapsule>,
reranker: &dyn crate::embeddings::Reranker,
floor: f32,
cap: usize,
) -> Vec<ContextCapsule> {
rerank_capsules_with_diagnostics(query, capsules, reranker, floor, cap).capsules
}
struct RerankOutcome {
capsules: Vec<ContextCapsule>,
best_raw_score: Option<f32>,
}
fn effective_rerank_policy_tier(capsule: &ContextCapsule) -> i8 {
if capsule.superseded_hint {
0
} else {
capsule.rerank_policy_tier
}
}
fn rerank_capsules_with_diagnostics(
query: &str,
capsules: Vec<ContextCapsule>,
reranker: &dyn crate::embeddings::Reranker,
floor: f32,
cap: usize,
) -> RerankOutcome {
if capsules.is_empty() {
return RerankOutcome {
capsules,
best_raw_score: None,
};
}
let docs: Vec<&str> = capsules.iter().map(|c| c.summary.as_str()).collect();
let scores = match reranker.rerank(query, &docs) {
Ok(s) if s.len() == docs.len() => s,
_ => {
let mut out = capsules;
if cap > 0 && out.len() > cap {
out.truncate(cap);
}
return RerankOutcome {
capsules: out,
best_raw_score: None,
};
}
};
let mut ranked: Vec<(ContextCapsule, f32)> = capsules
.into_iter()
.zip(scores)
.map(|(mut c, s)| {
let multiplier = if c.superseded_hint {
1.0
} else {
c.rerank_usefulness
.unwrap_or(1.0 + 0.5 * effective_rerank_policy_tier(&c) as f32)
};
c.score = apply_usefulness_boost(s, multiplier.clamp(0.5, 1.5))
* c.rerank_trust.unwrap_or(1.0).clamp(0.0, 1.0);
if c.superseded_hint {
c.score *= SUPERSESSION_PENALTY;
}
(c, s)
})
.collect();
ranked.sort_by(|a, b| b.0.score.total_cmp(&a.0.score));
ranked.retain(|(_, raw_score)| *raw_score >= floor);
let best_raw_score = ranked
.iter()
.map(|(_, raw_score)| *raw_score)
.max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
if cap > 0 && ranked.len() > cap {
ranked.truncate(cap);
}
RerankOutcome {
capsules: ranked.into_iter().map(|(c, _)| c).collect(),
best_raw_score,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn capsule(kind: &str, summary: &str) -> ContextCapsule {
ContextCapsule {
id: "c".into(),
kind: kind.into(),
summary: summary.into(),
token_estimate: 1,
expansion_handle: "memory:x".into(),
provenance: vec![],
confidence: 1.0,
freshness: 1.0,
relevance: 1.0,
scope_weight: 1.0,
score: 1.0,
superseded_hint: false,
rerank_policy_tier: 0,
claim_revision: None,
facts: vec![],
rerank_usefulness: None,
rerank_trust: None,
}
}
fn make_test_dir(tag: &str) -> std::path::PathBuf {
use std::time::{SystemTime, UNIX_EPOCH};
let ts = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.subsec_nanos())
.unwrap_or(0);
let dir = std::env::temp_dir().join(format!("kbrain_test_{tag}_{ts}"));
std::fs::create_dir_all(&dir).expect("create test dir");
dir
}
#[test]
fn capsule_matches_kind_reads_memory_summary_prefix() {
let mem = capsule("memory", "project:failure_pattern - linker not found");
assert!(capsule_matches_kind(&mem, "failure_pattern"));
assert!(!capsule_matches_kind(&mem, "command"));
let repo = capsule("repo_file", "src/lib.rs:command - run build");
assert!(capsule_matches_kind(&repo, "repo_file"));
assert!(!capsule_matches_kind(&repo, "command"));
}
#[test]
fn usefulness_multiplier_neutral_at_zero_uses() {
assert!((usefulness_multiplier(0.0, 0) - 1.0).abs() < f32::EPSILON);
assert!((usefulness_multiplier(5.0, 0) - 1.0).abs() < f32::EPSILON);
assert!((usefulness_multiplier(-5.0, 0) - 1.0).abs() < f32::EPSILON);
}
#[test]
fn usefulness_multiplier_blends_smoothly_in_transition() {
let one_use = usefulness_multiplier(1.0, 1);
assert!((one_use - 1.166_666_6).abs() < 1e-4, "got {one_use}");
let two_uses = usefulness_multiplier(2.0, 2);
assert!((two_uses - 1.333_333_4).abs() < 1e-4, "got {two_uses}");
let two_uses_bad = usefulness_multiplier(-2.0, 2);
assert!(
(two_uses_bad - 0.666_666_7).abs() < 1e-4,
"got {two_uses_bad}"
);
}
#[test]
fn usefulness_multiplier_maps_ratio_onto_envelope() {
assert!((usefulness_multiplier(5.0, 5) - 1.5).abs() < f32::EPSILON);
assert!((usefulness_multiplier(-5.0, 5) - 0.5).abs() < f32::EPSILON);
let mid = usefulness_multiplier(0.0, 6);
assert!((mid - 1.0).abs() < f32::EPSILON, "got {mid}");
let high = usefulness_multiplier(2.0, 4);
assert!((high - 1.25).abs() < f32::EPSILON, "got {high}");
let low = usefulness_multiplier(-2.0, 4);
assert!((low - 0.75).abs() < f32::EPSILON, "got {low}");
}
#[test]
fn usefulness_multiplier_clamps_to_envelope() {
assert!((usefulness_multiplier(100.0, 5) - 1.5).abs() < f32::EPSILON);
assert!((usefulness_multiplier(-100.0, 5) - 0.5).abs() < f32::EPSILON);
}
#[test]
fn boost_gain_is_capped_so_cited_junk_cannot_beat_relevant_uncited() {
let junk = apply_usefulness_boost(0.39, 1.5);
let true_match = apply_usefulness_boost(0.53, 1.0);
assert!(
junk < true_match,
"capped boost must preserve relevance order: junk {junk} vs match {true_match}"
);
assert!(junk <= 0.39 + USEFULNESS_BOOST_CAP + f32::EPSILON);
}
#[test]
fn boost_still_reorders_within_a_relevance_band() {
let proven = apply_usefulness_boost(0.85, 1.5);
let neutral = apply_usefulness_boost(0.90, 1.0);
assert!(
proven > neutral,
"capped boost must still reorder near-equals: proven {proven} vs neutral {neutral}"
);
}
#[test]
fn penalty_side_remains_multiplicative() {
let penalized = apply_usefulness_boost(0.8, 0.5);
assert!((penalized - 0.4).abs() < 1e-6);
}
#[test]
fn query_tokens_expands_build_class() {
let toks = query_tokens("Build the project from source");
assert!(toks.iter().any(|t| t == "build"));
assert!(toks.iter().any(|t| t == "shell_background"));
assert!(toks.iter().any(|t| t == "long_running"));
}
#[test]
fn query_tokens_expands_edit_class() {
let toks = query_tokens("Modify the config to fix the bug");
assert!(toks.iter().any(|t| t == "edit_file"));
assert!(toks.iter().any(|t| t == "apply_patch"));
}
#[test]
fn query_tokens_expands_search_class() {
let toks = query_tokens("Find all references to the symbol");
assert!(toks.iter().any(|t| t == "glob"));
assert!(toks.iter().any(|t| t == "search_files"));
}
#[test]
fn query_tokens_no_expansion_on_unrelated_query() {
let toks = query_tokens("hello world testing nothing");
assert!(toks.iter().any(|t| t == "hello"));
assert!(toks.iter().any(|t| t == "world"));
}
#[test]
fn jaccard_is_zero_for_disjoint_sets() {
let a: std::collections::HashSet<String> =
["foo", "bar"].iter().map(|s| s.to_string()).collect();
let b: std::collections::HashSet<String> =
["baz", "qux"].iter().map(|s| s.to_string()).collect();
assert!((jaccard(&a, &b) - 0.0).abs() < f32::EPSILON);
}
#[test]
fn jaccard_is_one_for_identical_sets() {
let a: std::collections::HashSet<String> =
["foo", "bar"].iter().map(|s| s.to_string()).collect();
let b = a.clone();
assert!((jaccard(&a, &b) - 1.0).abs() < f32::EPSILON);
}
#[test]
fn jaccard_partial_overlap() {
let a: std::collections::HashSet<String> = ["foo", "bar", "baz"]
.iter()
.map(|s| s.to_string())
.collect();
let b: std::collections::HashSet<String> =
["bar", "qux"].iter().map(|s| s.to_string()).collect();
assert!((jaccard(&a, &b) - 0.25).abs() < f32::EPSILON);
}
#[test]
fn summary_token_set_lowercases_and_filters_short() {
let set = summary_token_set("Build the Foo-bar project");
assert!(set.contains("build"));
assert!(set.contains("foo"));
assert!(set.contains("bar"));
assert!(set.contains("project"));
assert!(set.contains("the"));
}
fn insert_memory_with_embedding(
conn: &rusqlite::Connection,
memory_id: &str,
text: &str,
embedder: &dyn embeddings::Embedder,
) {
let normalized = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"
INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, embedding, embedding_model
)
VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
'2026-05-01T00:00:00Z', 0, 0.0, ?4, ?5)
",
rusqlite::params![
memory_id,
text,
normalized,
embeddings::encode_embedding(&embedder.embed(text).expect("embed test row")),
embedder.model_id(),
],
)
.expect("insert memory");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope) VALUES (?1, ?2, 'fact', 'global_user')",
rusqlite::params![memory_id, text],
)
.expect("insert fts row");
}
#[test]
fn hybrid_retrieval_uses_cosine_score_to_rerank() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let stub = embeddings::StubEmbedder::new();
insert_memory_with_embedding(&conn, "m_rg", "use ripgrep for code search", &stub);
insert_memory_with_embedding(
&conn,
"m_unrelated",
"cookie recipe with chocolate chips",
&stub,
);
let weights = kimetsu_core::config::BrokerWeights::default();
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "ripgrep search".to_string(),
budget_tokens: 4000,
..Default::default()
},
&[],
&stub,
)
.expect("retrieve");
let memory_handles: Vec<_> = bundle
.capsules
.iter()
.filter(|c| c.expansion_handle.starts_with("memory:"))
.collect();
assert!(
!memory_handles.is_empty(),
"at least one memory should surface"
);
assert_eq!(
memory_handles[0].expansion_handle,
"memory:m_rg",
"ripgrep memory should outrank the cookie recipe; ranked: {:?}",
memory_handles
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
}
#[test]
fn abstain_evidence_gate_skips_on_weak_absolute_evidence() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let stub = embeddings::StubEmbedder::new();
insert_memory_with_embedding(&conn, "m_rg", "use ripgrep for code search", &stub);
let weights = kimetsu_core::config::BrokerWeights::default();
let retrieve = |abstain: f32| {
retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "ripgrep search".to_string(),
budget_tokens: 4000,
abstain_evidence: abstain,
..Default::default()
},
&[],
&stub,
)
.expect("retrieve")
};
let open = retrieve(0.0);
assert!(!open.skipped, "gate off must not skip");
assert!(
open.top_abs_evidence > 0.0,
"a matching memory must report positive absolute evidence"
);
let above = retrieve(open.top_abs_evidence + ABSTAIN_BAND_WIDTH + 0.05);
assert!(
above.skipped,
"a floor a full band above the best evidence must hard-abstain (evidence {})",
open.top_abs_evidence
);
assert!(above.capsules.is_empty(), "skipped bundle injects nothing");
let in_band = retrieve(open.top_abs_evidence + 0.05);
assert!(
!in_band.skipped,
"an in-band bundle passes through for arbitration"
);
let below = retrieve((open.top_abs_evidence - 0.05).max(0.01));
assert!(!below.skipped, "a floor below the best evidence passes");
assert!(!below.capsules.is_empty());
}
#[test]
fn hybrid_retrieval_skips_cosine_on_model_id_mismatch() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let stub = embeddings::StubEmbedder::new();
insert_memory_with_embedding(&conn, "m_xref", "use ripgrep for code search", &stub);
conn.execute(
"UPDATE memories SET embedding_model = 'bge-small-en-v1.5' WHERE memory_id = 'm_xref'",
[],
)
.expect("force model_id mismatch");
let weights = kimetsu_core::config::BrokerWeights::default();
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "ripgrep search".to_string(),
budget_tokens: 4000,
..Default::default()
},
&[],
&stub,
)
.expect("retrieve");
assert!(
bundle
.capsules
.iter()
.any(|c| c.expansion_handle == "memory:m_xref"),
"cross-model row should still match lexically (cosine skipped, FTS works)"
);
}
#[test]
fn usefulness_decay_disabled_when_half_life_is_zero_or_negative() {
let ancient = "2021-01-01T00:00:00Z";
assert!((usefulness_decay(Some(ancient), ancient, 0.0) - 1.0).abs() < f32::EPSILON);
assert!((usefulness_decay(Some(ancient), ancient, -1.0) - 1.0).abs() < f32::EPSILON);
}
#[test]
fn usefulness_decay_returns_one_on_unparseable_timestamps() {
assert!(
(usefulness_decay(Some("not-a-date"), "also-not", 30.0) - 1.0).abs() < f32::EPSILON
);
}
#[test]
fn usefulness_decay_full_at_zero_age() {
let future = "2099-01-01T00:00:00Z";
let d = usefulness_decay(Some(future), future, 30.0);
assert!((d - 1.0).abs() < f32::EPSILON, "got {d}");
}
#[test]
fn usefulness_decay_follows_half_life_curve() {
let half_life = 10.0_f32;
let now = OffsetDateTime::now_utc();
let fmt = &time::format_description::well_known::Rfc3339;
let one_half_life_ago = (now - time::Duration::seconds((half_life * 86_400.0) as i64))
.format(fmt)
.expect("format");
let d1 = usefulness_decay(Some(&one_half_life_ago), &one_half_life_ago, half_life);
assert!(
(d1 - 0.5).abs() < 0.01,
"expected ~0.5 at one half-life, got {d1}"
);
let two_half_lives_ago = (now
- time::Duration::seconds((2.0 * half_life * 86_400.0) as i64))
.format(fmt)
.expect("format");
let d2 = usefulness_decay(Some(&two_half_lives_ago), &two_half_lives_ago, half_life);
assert!(
(d2 - 0.25).abs() < 0.01,
"expected ~0.25 at two half-lives, got {d2}"
);
}
#[test]
fn usefulness_decay_falls_back_to_created_at_when_last_useful_is_none() {
let now = OffsetDateTime::now_utc();
let fmt = &time::format_description::well_known::Rfc3339;
let one_day_ago = (now - time::Duration::seconds(86_400))
.format(fmt)
.expect("format");
let d = usefulness_decay(None, &one_day_ago, 30.0);
assert!(
(d - 0.977).abs() < 0.01,
"expected ~0.977 for 1-day-old created_at under 30d half-life, got {d}"
);
}
#[test]
fn aged_cited_memory_ranks_below_recently_cited_memory() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let now = OffsetDateTime::now_utc();
let fmt = &time::format_description::well_known::Rfc3339;
let one_day_ago = (now - time::Duration::seconds(86_400))
.format(fmt)
.expect("format");
let one_year_ago = (now - time::Duration::seconds(365 * 86_400))
.format(fmt)
.expect("format");
for (mid, last_useful) in [("m_recent", &one_day_ago), ("m_aged", &one_year_ago)] {
let text = "use ripgrep for code search";
let normalized = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"
INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, last_useful_at
)
VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
'2024-01-01T00:00:00Z', 5, 5.0, ?4)
",
rusqlite::params![mid, text, normalized, last_useful],
)
.expect("insert memory");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES (?1, ?2, 'fact', 'global_user')",
rusqlite::params![mid, text],
)
.expect("insert fts");
}
let weights = kimetsu_core::config::BrokerWeights::default();
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "ripgrep search".to_string(),
budget_tokens: 4000,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve");
let mem_order: Vec<&str> = bundle
.capsules
.iter()
.filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
.collect();
assert_eq!(
mem_order.first().copied(),
Some("m_recent"),
"recently-cited memory must rank first under decay; got order {mem_order:?}"
);
}
#[test]
fn aged_cited_memory_does_not_decay_when_half_life_is_zero() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let now = OffsetDateTime::now_utc();
let fmt = &time::format_description::well_known::Rfc3339;
let one_day_ago = (now - time::Duration::seconds(86_400))
.format(fmt)
.expect("format");
let one_year_ago = (now - time::Duration::seconds(365 * 86_400))
.format(fmt)
.expect("format");
for (mid, last_useful) in [("m_recent", &one_day_ago), ("m_aged", &one_year_ago)] {
let text = "use ripgrep for code search";
let normalized = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"
INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, last_useful_at
)
VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
'2024-01-01T00:00:00Z', 5, 5.0, ?4)
",
rusqlite::params![mid, text, normalized, last_useful],
)
.expect("insert memory");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES (?1, ?2, 'fact', 'global_user')",
rusqlite::params![mid, text],
)
.expect("insert fts");
}
let weights = kimetsu_core::config::BrokerWeights {
decay_half_life_days: 0.0,
..Default::default()
};
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "ripgrep search".to_string(),
budget_tokens: 4000,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve");
let scores: Vec<(String, f32)> = bundle
.capsules
.iter()
.filter_map(|c| {
c.expansion_handle
.strip_prefix("memory:")
.map(|id| (id.to_string(), c.score))
})
.collect();
assert_eq!(scores.len(), 2, "both memories should surface");
let recent_score = scores
.iter()
.find(|(id, _)| id == "m_recent")
.map(|(_, s)| *s)
.expect("m_recent present");
let aged_score = scores
.iter()
.find(|(id, _)| id == "m_aged")
.map(|(_, s)| *s)
.expect("m_aged present");
assert!(
(recent_score - aged_score).abs() < 1e-4,
"with decay disabled the two memories should tie on score: recent={recent_score} aged={aged_score}"
);
}
#[test]
fn hybrid_retrieval_with_noop_embedder_is_lexical_only() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let stub = embeddings::StubEmbedder::new();
insert_memory_with_embedding(&conn, "m_a", "use ripgrep", &stub);
insert_memory_with_embedding(&conn, "m_b", "use ripgrep too", &stub);
let weights = kimetsu_core::config::BrokerWeights::default();
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "ripgrep".to_string(),
budget_tokens: 4000,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve");
let count = bundle
.capsules
.iter()
.filter(|c| c.expansion_handle.starts_with("memory:"))
.count();
assert_eq!(count, 2, "both memories should surface via FTS");
}
#[cfg(feature = "embeddings")]
#[test]
fn ann_finds_semantic_match_fts_misses() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
struct OracleEmbedder;
impl embeddings::Embedder for OracleEmbedder {
fn embed(&self, _text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
Ok(vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
}
fn model_id(&self) -> &str {
"oracle-d8"
}
fn dim(&self) -> usize {
8
}
}
let model_id = "oracle-d8";
let sem_vec = embeddings::encode_embedding(&[1.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
let sem_text = "cookie recipe chocolate";
let sem_norm = kimetsu_core::memory::normalize_memory_text(sem_text);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, embedding, embedding_model
)
VALUES ('m_semantic', 'global_user', 'fact', ?1, ?2, 1.0, NULL, '{}',
'2026-01-01T00:00:00Z', 0, 0.0, ?3, ?4)",
rusqlite::params![sem_text, sem_norm, sem_vec, model_id],
)
.expect("insert m_semantic");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES ('m_semantic', ?1, 'fact', 'global_user')",
rusqlite::params![sem_text],
)
.expect("insert m_semantic fts");
let decoy_vec = embeddings::encode_embedding(&[0.0f32, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
let decoy_text = "git rebase squash commits";
let decoy_norm = kimetsu_core::memory::normalize_memory_text(decoy_text);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, embedding, embedding_model
)
VALUES ('m_decoy', 'global_user', 'fact', ?1, ?2, 1.0, NULL, '{}',
'2026-01-01T00:00:00Z', 0, 0.0, ?3, ?4)",
rusqlite::params![decoy_text, decoy_norm, decoy_vec, model_id],
)
.expect("insert m_decoy");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES ('m_decoy', ?1, 'fact', 'global_user')",
rusqlite::params![decoy_text],
)
.expect("insert m_decoy fts");
let fts_hits: i64 = conn
.query_row(
"SELECT COUNT(*) FROM memories_fts \
WHERE memories_fts MATCH 'phosphorescent bioluminescent'",
[],
|r| r.get(0),
)
.unwrap_or(0);
assert_eq!(
fts_hits, 0,
"sanity: query tokens must not appear in any memory text"
);
let weights = kimetsu_core::config::BrokerWeights::default();
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "phosphorescent bioluminescent organism".to_string(),
budget_tokens: 4000,
..Default::default()
},
&[],
&OracleEmbedder,
)
.expect("retrieve");
let handles: Vec<&str> = bundle
.capsules
.iter()
.filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
.collect();
assert!(
handles.contains(&"m_semantic"),
"ANN must surface m_semantic (cosine=1 with oracle query) even though \
FTS found nothing; got handles: {handles:?}"
);
}
#[cfg(feature = "embeddings")]
#[test]
fn dedup_memory_matched_by_fts_and_ann_appears_once() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let stub = embeddings::StubEmbedder::new();
insert_memory_with_embedding(&conn, "m_both", "use ripgrep for fast search", &stub);
let weights = kimetsu_core::config::BrokerWeights::default();
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "ripgrep".to_string(),
budget_tokens: 4000,
..Default::default()
},
&[],
&stub,
)
.expect("retrieve");
let count = bundle
.capsules
.iter()
.filter(|c| c.expansion_handle == "memory:m_both")
.count();
assert_eq!(
count,
1,
"m_both (matched by both FTS and ANN) must appear exactly once; \
bundle: {:?}",
bundle
.capsules
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
}
#[cfg(feature = "embeddings")]
#[test]
fn embedding_mmr_collapses_paraphrases_but_jaccard_does_not() {
struct OracleEmbedder;
impl embeddings::Embedder for OracleEmbedder {
fn embed(&self, _text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
let mut v = vec![0.0f32; 8];
v[0] = 1.0;
Ok(v)
}
fn model_id(&self) -> &str {
"oracle-d8"
}
fn dim(&self) -> usize {
8
}
}
let oracle = OracleEmbedder;
let weights = kimetsu_core::config::BrokerWeights::default();
let m_rg1_text = "prefer ripgrep for searching source code";
let m_rg2_text = "rg is the fastest way to locate patterns";
let conn = rusqlite::Connection::open_in_memory().expect("in-memory");
crate::schema::initialize(&conn).expect("init schema");
insert_memory_with_embedding(&conn, "m_rg1", m_rg1_text, &oracle);
insert_memory_with_embedding(&conn, "m_rg2", m_rg2_text, &oracle);
let bundle_embedding = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "search source patterns".to_string(),
budget_tokens: 20_000,
max_capsules: 1, ..Default::default()
},
&[],
&oracle,
)
.expect("retrieve with oracle embedder");
let emb_in_capsules = bundle_embedding
.capsules
.iter()
.filter(|c| {
c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
})
.count();
assert_eq!(
emb_in_capsules,
1,
"embedding-MMR must collapse cosine=1.0 paraphrases: with max_capsules=1 \
only ONE should be included; capsule handles: {:?}; excluded: {:?}",
bundle_embedding
.capsules
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>(),
bundle_embedding
.excluded
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
let emb_in_excluded = bundle_embedding
.excluded
.iter()
.filter(|c| {
c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
})
.count();
assert_eq!(
emb_in_excluded,
1,
"the second near-duplicate must be in excluded under embedding-MMR; \
excluded handles: {:?}",
bundle_embedding
.excluded
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
let conn2 = rusqlite::Connection::open_in_memory().expect("in-memory 2");
crate::schema::initialize(&conn2).expect("init schema 2");
insert_memory_with_embedding(&conn2, "m_rg1", m_rg1_text, &oracle);
insert_memory_with_embedding(&conn2, "m_rg2", m_rg2_text, &oracle);
let bundle_lean = retrieve_context_with_embedder(
&conn2,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "search source patterns".to_string(),
budget_tokens: 20_000,
max_capsules: 2, ..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve with NoopEmbedder");
let lean_in_capsules = bundle_lean
.capsules
.iter()
.filter(|c| {
c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
})
.count();
assert_eq!(
lean_in_capsules,
2,
"Jaccard-only path must NOT collapse the two paraphrases (different words, \
low token overlap → both survive MMR with max_capsules=2); capsule handles: {:?}",
bundle_lean
.capsules
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
}
#[test]
fn content_tokens_strips_stopwords_keeps_topical_words() {
let got = content_tokens("Tell me about kimetsu, what's the idea of the repo");
assert_eq!(got, vec!["kimetsu", "idea", "repo"]);
}
#[test]
fn light_stem_strips_one_inflection_suffix() {
assert_eq!(light_stem("benchmarked"), "benchmark");
assert_eq!(light_stem("benchmarking"), "benchmark");
assert_eq!(light_stem("repos"), "repo");
assert_eq!(light_stem("does"), "does");
assert_eq!(light_stem("toml"), "toml");
}
#[test]
fn stemmed_query_matches_inflected_corpus_through_floor() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let insert = |id: &str, text: &str| {
let norm = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, embedding, embedding_model
)
VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
'2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
rusqlite::params![id, text, norm],
)
.expect("insert memory");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES (?1, ?2, 'fact', 'global_user')",
rusqlite::params![id, text],
)
.expect("insert fts");
};
insert(
"m_bench",
"kimetsu benchmark runs go through the kbench binary and the Terminal-Bench driver",
);
insert(
"m_doctor",
"kimetsu doctor version-skew check parses process start times on Windows via CIM",
);
insert(
"m_gc",
"kimetsu runs auto-GC on run creation; keep the env guard at the trigger site",
);
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&kimetsu_core::config::BrokerWeights::default(),
ContextRequest {
stage: "localization".to_string(),
query: "Can you find out how kimetsu is benchmarked?".to_string(),
budget_tokens: 2000,
max_capsules: 2,
min_lexical_coverage: 0.5,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve");
let handles: Vec<_> = bundle
.capsules
.iter()
.map(|c| c.expansion_handle.as_str())
.collect();
assert!(
handles.contains(&"memory:m_bench"),
"stemmed 'benchmarked' must surface the benchmark memory; got {handles:?}"
);
assert!(
!handles.contains(&"memory:m_doctor") && !handles.contains(&"memory:m_gc"),
"off-topic memories sharing only 'kimetsu' must stay below the floor; got {handles:?}"
);
}
#[test]
fn weighted_coverage_ignores_zero_idf_tokens() {
let content = vec![
"kimetsu".to_string(),
"idea".to_string(),
"repo".to_string(),
];
let mut idf = HashMap::new();
idf.insert("kimetsu".to_string(), 0.0);
idf.insert("idea".to_string(), 1.386);
idf.insert("repo".to_string(), 0.693);
let cov = weighted_coverage(
&content,
&idf,
"global:fact - the git repo and kimetsu brain",
);
assert!((cov - 0.333).abs() < 0.01, "got {cov}");
let cov_topical =
weighted_coverage(&content, &idf, "global:fact - the core idea of kimetsu");
assert!(cov_topical > 0.6, "got {cov_topical}");
}
#[test]
fn lexical_floor_drops_offtopic_memories_sharing_project_name() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let insert = |id: &str, text: &str| {
let norm = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, embedding, embedding_model
)
VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
'2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
rusqlite::params![id, text, norm],
)
.expect("insert memory");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES (?1, ?2, 'fact', 'global_user')",
rusqlite::params![id, text],
)
.expect("insert fts");
};
insert(
"m1",
"When implementing a setup command that calls init_project, tests must call \
git_init_boundary before setup_cmd so ProjectPaths discover resolves to the temp \
dir instead of climbing to the real parent git repo including the user brain at kimetsu",
);
insert(
"m2",
"A member crate with default embeddings silently turned embeddings on for the entire \
cargo test workspace build graph because cargo unifies features; kimetsu-chat \
retrieval tests failed",
);
insert(
"m3",
"In toml 0.9 use toml from_str to parse a TOML document into a Value not str parse; \
implementing config get and set in kimetsu-cli",
);
let query = "Tell me about kimetsu, what's the idea of the repo".to_string();
let weights = kimetsu_core::config::BrokerWeights::default();
let handles = |bundle: &ContextBundle| {
bundle
.capsules
.iter()
.map(|c| c.expansion_handle.clone())
.collect::<Vec<_>>()
};
let no_floor = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: query.clone(),
budget_tokens: 2000,
max_capsules: 8,
min_lexical_coverage: 0.0,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve without floor");
let before = handles(&no_floor);
assert!(
before.contains(&"memory:m2".to_string()) && before.contains(&"memory:m3".to_string()),
"sanity: without the floor the pure-project-name memories should surface; got {before:?}"
);
let floored = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query,
budget_tokens: 2000,
max_capsules: 8,
min_lexical_coverage: 0.5,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve with floor");
let after = handles(&floored);
assert!(
!after.contains(&"memory:m2".to_string()) && !after.contains(&"memory:m3".to_string()),
"the lexical floor must drop memories whose only match is the corpus-ubiquitous \
project name; surviving: {after:?}"
);
}
#[test]
fn lexical_floor_keeps_ontopic_memory() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let insert = |id: &str, text: &str| {
let norm = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, embedding, embedding_model
)
VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
'2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
rusqlite::params![id, text, norm],
)
.expect("insert memory");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES (?1, ?2, 'fact', 'global_user')",
rusqlite::params![id, text],
)
.expect("insert fts");
};
insert(
"d1",
"The distiller runs at session end and harvests durable lessons from the transcript",
);
insert(
"n1",
"Unrelated note about git rebase and squashing commits",
);
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&kimetsu_core::config::BrokerWeights::default(),
ContextRequest {
stage: "localization".to_string(),
query: "how does the distiller work".to_string(),
budget_tokens: 2000,
min_lexical_coverage: 0.5,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve");
assert!(
bundle
.capsules
.iter()
.any(|c| c.expansion_handle == "memory:d1"),
"on-topic memory covering the rare query word must survive the floor; got: {:?}",
bundle
.capsules
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
}
#[cfg(feature = "embeddings")]
#[test]
fn min_semantic_score_floor_drops_off_topic_queries() {
struct DirectionalEmbedder {
marker: &'static str,
}
impl embeddings::Embedder for DirectionalEmbedder {
fn embed(&self, text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
let mut v = vec![0.0f32; 8];
if text.contains(self.marker) {
v[0] = 1.0;
} else {
v[1] = 1.0;
}
Ok(v)
}
fn model_id(&self) -> &str {
"directional-d8"
}
fn dim(&self) -> usize {
8
}
}
let emb = DirectionalEmbedder { marker: "TOPIC_A" };
let conn = rusqlite::Connection::open_in_memory().expect("in-memory");
crate::schema::initialize(&conn).expect("init schema");
insert_memory_with_embedding(&conn, "m_b", "cookie recipe chocolate baking TOPIC_B", &emb);
let weights = kimetsu_core::config::BrokerWeights::default();
let bundle_off = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "TOPIC_A unrelated phosphorescent".to_string(),
budget_tokens: 4000,
min_semantic_score: 0.1, ..Default::default()
},
&[],
&emb,
)
.expect("retrieve off-topic");
assert!(
bundle_off.capsules.is_empty(),
"off-topic query (cosine=0 < floor=0.1) must produce zero capsules; \
got: {:?}",
bundle_off
.capsules
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
let conn2 = rusqlite::Connection::open_in_memory().expect("in-memory 2");
crate::schema::initialize(&conn2).expect("init schema 2");
insert_memory_with_embedding(
&conn2,
"m_b2",
"cookie recipe chocolate TOPIC_B baking"
.to_string()
.as_str(),
&emb,
);
let bundle_on = retrieve_context_with_embedder(
&conn2,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "cookie chocolate TOPIC_B".to_string(),
budget_tokens: 4000,
min_semantic_score: 0.1,
..Default::default()
},
&[],
&emb,
)
.expect("retrieve on-topic");
assert!(
bundle_on
.capsules
.iter()
.any(|c| c.expansion_handle == "memory:m_b2"),
"on-topic query (cosine=1.0 ≥ floor) must surface m_b2; \
got capsules: {:?}",
bundle_on
.capsules
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
let conn3 = rusqlite::Connection::open_in_memory().expect("in-memory 3");
crate::schema::initialize(&conn3).expect("init schema 3");
insert_memory_with_embedding(
&conn3,
"m_b3",
"cookie chocolate TOPIC_B recipe".to_string().as_str(),
&emb,
);
let bundle_noop_floor = retrieve_context_with_embedder(
&conn3,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "cookie chocolate TOPIC_A".to_string(),
budget_tokens: 4000,
min_semantic_score: 0.0, ..Default::default()
},
&[],
&emb,
)
.expect("retrieve noop floor");
assert!(
bundle_noop_floor
.capsules
.iter()
.any(|c| c.expansion_handle == "memory:m_b3"),
"with floor=0.0 (disabled), off-topic-cosine memory must still surface via FTS; \
got: {:?}",
bundle_noop_floor
.capsules
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
}
#[cfg(feature = "embeddings")]
#[test]
fn d1f_token_economy_fewer_capsules_signal_preserved() {
struct OracleTopicEmbedder;
impl embeddings::Embedder for OracleTopicEmbedder {
fn embed(&self, text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
let mut v = vec![0.0f32; 8];
if text.contains("TOPIC_A") {
v[0] = 1.0; } else {
v[1] = 1.0; }
Ok(v)
}
fn model_id(&self) -> &str {
"oracle-topic-d8"
}
fn dim(&self) -> usize {
8
}
}
let oracle = OracleTopicEmbedder;
let setup = |conn: &rusqlite::Connection| {
for (mid, text) in [
("m_dup1", "TOPIC_A prefer ripgrep for searching"),
("m_dup2", "TOPIC_A rg is the fastest searcher"),
("m_dup3", "TOPIC_A use rg tool to find patterns"),
(
"m_relevant",
"TOPIC_A critical lesson about search performance",
),
("m_noise1", "chocolate cookie baking TOPIC_B recipe"),
("m_noise2", "gardening tulip planting TOPIC_B spring"),
] {
insert_memory_with_embedding(conn, mid, text, &oracle);
}
};
let weights = kimetsu_core::config::BrokerWeights::default();
let conn_lean = rusqlite::Connection::open_in_memory().expect("in-memory lean");
crate::schema::initialize(&conn_lean).expect("init schema lean");
setup(&conn_lean);
let bundle_lean = retrieve_context_with_embedder(
&conn_lean,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "TOPIC_A search performance".to_string(),
budget_tokens: 20_000,
min_semantic_score: 0.0, ..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve lean");
let lean_count = bundle_lean
.capsules
.iter()
.filter(|c| c.expansion_handle.starts_with("memory:"))
.count();
let conn_emb = rusqlite::Connection::open_in_memory().expect("in-memory emb");
crate::schema::initialize(&conn_emb).expect("init schema emb");
setup(&conn_emb);
let bundle_emb = retrieve_context_with_embedder(
&conn_emb,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "TOPIC_A search performance".to_string(),
budget_tokens: 20_000,
min_semantic_score: 0.5, ..Default::default()
},
&[],
&oracle,
)
.expect("retrieve with embeddings");
let emb_count = bundle_emb
.capsules
.iter()
.filter(|c| c.expansion_handle.starts_with("memory:"))
.count();
assert!(
emb_count < lean_count,
"D1e must reduce capsule count: embedding path {emb_count} must be \
< lean path {lean_count}. Embedding capsules: {:?}",
bundle_emb
.capsules
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
assert!(
bundle_emb
.capsules
.iter()
.any(|c| c.expansion_handle == "memory:m_relevant"),
"m_relevant must survive D1e selection (signal preserved); \
embedding capsules: {:?}",
bundle_emb
.capsules
.iter()
.map(|c| &c.expansion_handle)
.collect::<Vec<_>>()
);
let lean_tokens: u32 = bundle_lean.capsules.iter().map(|c| c.token_estimate).sum();
let emb_tokens: u32 = bundle_emb.capsules.iter().map(|c| c.token_estimate).sum();
assert!(
emb_tokens < lean_tokens,
"D1e must reduce token usage: emb={emb_tokens} must be < lean={lean_tokens}"
);
}
#[test]
fn lean_noop_embedder_uses_fts_then_recency_unchanged() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
for (mid, text) in [
("m_x", "use git rebase to clean history"),
("m_y", "grep finds text quickly"),
] {
let normalized = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score
)
VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
'2026-01-01T00:00:00Z', 0, 0.0)",
rusqlite::params![mid, text, normalized],
)
.expect("insert");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope) VALUES (?1, ?2, 'fact', 'global_user')",
rusqlite::params![mid, text],
)
.expect("insert fts");
}
let weights = kimetsu_core::config::BrokerWeights::default();
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: "grep text".to_string(),
budget_tokens: 4000,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve with NoopEmbedder must not panic");
let handles: Vec<&str> = bundle
.capsules
.iter()
.filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
.collect();
assert!(
handles.contains(&"m_y"),
"m_y must surface via FTS on lean path; got {handles:?}"
);
}
#[test]
fn classify_task_maps_each_kind_deterministically() {
assert_eq!(
classify_task("fix the panic in the parser"),
TaskKind::Debug,
"contains 'fix' and 'panic'"
);
assert_eq!(
classify_task("there is a crash in auth when calling login"),
TaskKind::Debug,
"contains 'crash'"
);
assert_eq!(
classify_task("debug the failing test"),
TaskKind::Debug,
"contains 'debug' and 'fail'"
);
assert_eq!(
classify_task("investigate why retrieval is slow"),
TaskKind::Investigation,
"contains 'investigate' and 'why'"
);
assert_eq!(
classify_task("analyze the root cause of the latency"),
TaskKind::Investigation,
"contains 'analyze' and 'root cause'"
);
assert_eq!(
classify_task("refactor the auth module"),
TaskKind::Refactor,
"contains 'refactor'"
);
assert_eq!(
classify_task("rename the config struct"),
TaskKind::Refactor,
"contains 'rename'"
);
assert_eq!(
classify_task("simplify the retry handling logic"),
TaskKind::Refactor,
"contains 'simplify'"
);
assert_eq!(
classify_task("document the API endpoints"),
TaskKind::Docs,
"contains 'document'"
);
assert_eq!(
classify_task("update the readme with new instructions"),
TaskKind::Docs,
"contains 'readme'"
);
assert_eq!(
classify_task("add a docstring to the main function"),
TaskKind::Docs,
"contains 'docstring'"
);
assert_eq!(
classify_task("add a dark mode toggle"),
TaskKind::Feature,
"no debug/refactor/docs/investigate keyword"
);
assert_eq!(
classify_task("implement the new caching layer"),
TaskKind::Feature,
"no debug/refactor/docs/investigate keyword"
);
assert_eq!(
classify_task("build the export pipeline"),
TaskKind::Feature,
"no debug/refactor/docs/investigate keyword"
);
}
#[test]
fn classify_task_respects_precedence_order() {
assert_eq!(
classify_task("fix and refactor the login module"),
TaskKind::Debug,
"Debug > Refactor"
);
assert_eq!(
classify_task("investigate and refactor the cache layer"),
TaskKind::Investigation,
"Investigation > Refactor"
);
assert_eq!(
classify_task("investigate the docs and document the API"),
TaskKind::Investigation,
"Investigation > Docs"
);
assert_eq!(
classify_task("refactor and add docs"),
TaskKind::Refactor,
"Refactor > Docs"
);
assert_eq!(
classify_task("fix the bug and investigate the regression"),
TaskKind::Debug,
"Debug > Investigation"
);
}
fn two_kinds_one_strong() -> Vec<Candidate> {
let mk = |kind: &str, raw: f32| Candidate {
capsule: ContextCapsule {
id: format!("{kind}-1"),
kind: kind.to_string(),
summary: String::new(),
token_estimate: 0,
expansion_handle: String::new(),
provenance: Vec::new(),
confidence: 0.0,
freshness: 0.0,
relevance: 0.0,
scope_weight: 0.0,
score: 0.0,
superseded_hint: false,
rerank_policy_tier: 0,
claim_revision: None,
facts: vec![],
rerank_usefulness: None,
rerank_trust: None,
},
raw_relevance: raw,
embedding: None,
cosine: None,
created_at: None,
};
vec![mk("memory", 0.9), mk("repo_file", 0.1)]
}
#[test]
fn per_kind_normalization_flatters_the_best_of_a_weak_kind() {
let mut candidates = two_kinds_one_strong();
let weights = StageWeights {
relevance: 1.0,
confidence: 0.0,
freshness: 0.0,
scope: 0.0,
};
normalize_and_score(&mut candidates, weights, Normalization::PerKind);
assert!((candidates[0].capsule.relevance - 1.0).abs() < 1e-6);
assert!(
(candidates[1].capsule.relevance - 1.0).abs() < 1e-6,
"per-kind gives the lone weak repo_file relevance 1.0, got {}",
candidates[1].capsule.relevance
);
}
#[test]
fn global_normalization_keeps_relevance_comparable_across_kinds() {
let mut candidates = two_kinds_one_strong();
let weights = StageWeights {
relevance: 1.0,
confidence: 0.0,
freshness: 0.0,
scope: 0.0,
};
normalize_and_score(&mut candidates, weights, Normalization::Global);
assert!((candidates[0].capsule.relevance - 1.0).abs() < 1e-6);
let weak = candidates[1].capsule.relevance;
assert!(
(weak - (0.1 / 0.9)).abs() < 1e-6,
"global normalizes against the single max, got {weak}"
);
assert!(weak < candidates[0].capsule.relevance);
}
fn superseding_candidate(
id: &str,
embedding: Vec<f32>,
created_at: &str,
score: f32,
) -> Candidate {
Candidate {
capsule: ContextCapsule {
id: id.to_string(),
kind: "memory".to_string(),
summary: id.to_string(),
token_estimate: 0,
expansion_handle: format!("memory:{id}"),
provenance: Vec::new(),
confidence: 0.0,
freshness: 0.0,
relevance: 0.0,
scope_weight: 0.0,
score,
superseded_hint: false,
rerank_policy_tier: 0,
claim_revision: None,
facts: vec![],
rerank_usefulness: None,
rerank_trust: None,
},
raw_relevance: score,
embedding: Some(embedding),
cosine: Some(score),
created_at: Some(created_at.to_string()),
}
}
#[test]
fn supersession_penalizes_the_older_near_duplicate() {
let mut candidates = vec![
superseding_candidate("old", vec![1.0, 0.0], "2026-08-01T10:00:00Z", 0.94),
superseding_candidate("new", vec![0.99, 0.14], "2026-08-01T10:10:00Z", 0.87),
];
apply_supersession_penalty(&mut candidates);
let old_score = candidates[0].capsule.score;
let new_score = candidates[1].capsule.score;
assert!(
(old_score - 0.94 * SUPERSESSION_PENALTY).abs() < 1e-6,
"older twin must carry the penalty, got {old_score}"
);
assert!((new_score - 0.87).abs() < 1e-6, "newer twin untouched");
assert!(
new_score > old_score,
"the update must now outrank the incumbent"
);
}
#[test]
fn supersession_ignores_distinct_memories_and_applies_once() {
let mut candidates = vec![
superseding_candidate("old", vec![1.0, 0.0], "2026-08-01T10:00:00Z", 0.90),
superseding_candidate("other", vec![0.0, 1.0], "2026-08-02T10:00:00Z", 0.80),
superseding_candidate("new1", vec![0.99, 0.14], "2026-08-03T10:00:00Z", 0.70),
superseding_candidate("new2", vec![0.98, 0.19], "2026-08-04T10:00:00Z", 0.60),
];
apply_supersession_penalty(&mut candidates);
assert!(
(candidates[0].capsule.score - 0.90 * SUPERSESSION_PENALTY).abs() < 1e-6,
"penalty applies exactly once, got {}",
candidates[0].capsule.score
);
assert!(
(candidates[1].capsule.score - 0.80).abs() < 1e-6,
"orthogonal memory untouched"
);
assert!(
(candidates[2].capsule.score - 0.70 * SUPERSESSION_PENALTY).abs() < 1e-6,
"a middle sibling is old relative to a newer one"
);
assert!(
(candidates[3].capsule.score - 0.60).abs() < 1e-6,
"newest untouched"
);
}
#[test]
fn supersession_is_inert_without_embeddings_or_timestamps() {
let mut no_embedding = vec![
Candidate {
embedding: None,
..superseding_candidate("a", vec![], "2026-08-01T10:00:00Z", 0.9)
},
Candidate {
embedding: None,
..superseding_candidate("b", vec![], "2026-08-02T10:00:00Z", 0.8)
},
];
apply_supersession_penalty(&mut no_embedding);
assert!((no_embedding[0].capsule.score - 0.9).abs() < 1e-6);
let mut bad_ts = vec![
superseding_candidate("a", vec![1.0, 0.0], "not-a-date", 0.9),
superseding_candidate("b", vec![1.0, 0.0], "2026-08-02T10:00:00Z", 0.8),
];
apply_supersession_penalty(&mut bad_ts);
assert!(
(bad_ts[0].capsule.score - 0.9).abs() < 1e-6,
"unparseable ts skipped"
);
assert!((bad_ts[1].capsule.score - 0.8).abs() < 1e-6);
let mut same_ts = vec![
superseding_candidate("a", vec![1.0, 0.0], "2026-08-01T10:00:00Z", 0.9),
superseding_candidate("b", vec![1.0, 0.0], "2026-08-01T10:00:00Z", 0.8),
];
apply_supersession_penalty(&mut same_ts);
assert!((same_ts[0].capsule.score - 0.9).abs() < 1e-6);
assert!((same_ts[1].capsule.score - 0.8).abs() < 1e-6);
let mut batch = vec![
superseding_candidate("a", vec![1.0, 0.0], "2026-08-01T10:00:00.100Z", 0.9),
superseding_candidate("b", vec![1.0, 0.0], "2026-08-01T10:00:00.900Z", 0.8),
];
apply_supersession_penalty(&mut batch);
assert!(
(batch[0].capsule.score - 0.9).abs() < 1e-6,
"millisecond-apart co-writes must not be penalized"
);
assert!((batch[1].capsule.score - 0.8).abs() < 1e-6);
let mut explicit_update = vec![
superseding_candidate(
"the project uses spaces (switched from tabs)",
vec![1.0, 0.0],
"2026-08-01T10:00:00.100Z",
0.9,
),
superseding_candidate(
"the project uses tabs for indentation",
vec![0.83, 0.557_8],
"2026-08-01T10:00:00.900Z",
0.8,
),
];
apply_supersession_penalty(&mut explicit_update);
assert!((explicit_update[0].capsule.score - 0.9).abs() < 1e-6);
assert!(
(explicit_update[1].capsule.score - 0.8 * SUPERSESSION_PENALTY).abs() < 1e-6,
"the unmarked incumbent must lose to the explicit correction"
);
assert!(explicit_update[1].capsule.superseded_hint);
let mut rewritten_update = vec![
superseding_candidate(
"as of v2 the preferred kimetsu embedder is jina, replacing bge",
vec![1.0, 0.0],
"2026-08-01T10:00:00.100Z",
0.9,
),
superseding_candidate(
"the recommended kimetsu embedder for retrieval is bge",
vec![0.71, 0.704_2],
"2026-08-01T10:00:00.900Z",
0.8,
),
];
apply_supersession_penalty(&mut rewritten_update);
assert!((rewritten_update[0].capsule.score - 0.9).abs() < 1e-6);
assert!((rewritten_update[1].capsule.score - 0.8 * SUPERSESSION_PENALTY).abs() < 1e-6);
let mut historical_query = vec![
Candidate {
cosine: Some(0.70),
..superseding_candidate(
"as of v2 the preferred kimetsu embedder is jina, replacing bge",
vec![1.0, 0.0],
"2026-08-01T10:00:00.100Z",
0.9,
)
},
Candidate {
cosine: Some(0.90),
..superseding_candidate(
"the recommended kimetsu embedder for retrieval is bge",
vec![0.71, 0.704_2],
"2026-08-01T10:00:00.900Z",
0.8,
)
},
];
apply_supersession_penalty(&mut historical_query);
assert!((historical_query[0].capsule.score - 0.9).abs() < 1e-6);
assert!((historical_query[1].capsule.score - 0.8).abs() < 1e-6);
let mut too_distant = vec![
superseding_candidate(
"the project now uses spaces",
vec![1.0, 0.0],
"2026-08-01T10:00:00.100Z",
0.9,
),
superseding_candidate(
"database backup retention is seven days",
vec![0.81, 0.586_4],
"2026-08-01T10:00:00.900Z",
0.8,
),
];
apply_supersession_penalty(&mut too_distant);
assert!((too_distant[0].capsule.score - 0.9).abs() < 1e-6);
assert!((too_distant[1].capsule.score - 0.8).abs() < 1e-6);
let mut ambiguous = vec![
superseding_candidate(
"the setting is now cheap_model",
vec![1.0, 0.0],
"2026-08-01T10:00:00.100Z",
0.9,
),
superseding_candidate(
"as of v2 the setting is cheap_model",
vec![1.0, 0.0],
"2026-08-01T10:00:00.900Z",
0.8,
),
];
apply_supersession_penalty(&mut ambiguous);
assert!((ambiguous[0].capsule.score - 0.9).abs() < 1e-6);
assert!((ambiguous[1].capsule.score - 0.8).abs() < 1e-6);
}
#[test]
fn unknown_normalization_falls_back_to_per_kind() {
assert_eq!(Normalization::from_config(""), Normalization::PerKind);
assert_eq!(
Normalization::from_config("per_kind"),
Normalization::PerKind
);
assert_eq!(
Normalization::from_config("nonsense"),
Normalization::PerKind
);
assert_eq!(Normalization::from_config("global"), Normalization::Global);
assert_eq!(
Normalization::from_config(" GLOBAL "),
Normalization::Global
);
}
#[test]
fn weights_for_task_kind_renormalizes_to_unit_sum() {
let base = StageWeights {
relevance: 0.50,
confidence: 0.20,
freshness: 0.20,
scope: 0.10,
};
let original_sum = base.relevance + base.confidence + base.freshness + base.scope;
for kind in [
TaskKind::Debug,
TaskKind::Refactor,
TaskKind::Investigation,
TaskKind::Docs,
] {
let w = weights_for_task_kind(base.clone(), kind);
let new_sum = w.relevance + w.confidence + w.freshness + w.scope;
assert!(
(new_sum - original_sum).abs() < 1e-4,
"weights_for_task_kind({kind:?}) sum {new_sum} differs from {original_sum}"
);
}
}
#[test]
fn weights_for_task_kind_feature_is_unchanged() {
let base = StageWeights {
relevance: 0.40,
confidence: 0.30,
freshness: 0.20,
scope: 0.10,
};
let w = weights_for_task_kind(base.clone(), TaskKind::Feature);
assert!((w.relevance - base.relevance).abs() < f32::EPSILON);
assert!((w.confidence - base.confidence).abs() < f32::EPSILON);
assert!((w.freshness - base.freshness).abs() < f32::EPSILON);
assert!((w.scope - base.scope).abs() < f32::EPSILON);
}
#[test]
fn weights_for_task_kind_debug_up_freshness_fraction() {
let base = StageWeights {
relevance: 0.50,
confidence: 0.20,
freshness: 0.20,
scope: 0.10,
};
let debug_w = weights_for_task_kind(base.clone(), TaskKind::Debug);
assert!(
debug_w.freshness > base.freshness,
"Debug must increase freshness fraction: {debug_w:?}"
);
}
#[test]
fn weights_for_task_kind_refactor_up_scope_fraction() {
let base = StageWeights {
relevance: 0.50,
confidence: 0.20,
freshness: 0.20,
scope: 0.10,
};
let refactor_w = weights_for_task_kind(base.clone(), TaskKind::Refactor);
assert!(
refactor_w.scope > base.scope,
"Refactor must increase scope fraction: {refactor_w:?}"
);
}
#[test]
fn task_kind_feature_is_retrieval_neutral() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
for (mid, db_kind, text) in [
("m1", "failure_pattern", "linker not found error in build"),
("m2", "convention", "use snake_case for all identifiers"),
("m3", "fact", "the cache is invalidated on every deploy"),
] {
let normalized = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score
)
VALUES (?1, 'project', ?2, ?3, ?4, 1.0, NULL, '{}',
'2026-01-01T00:00:00Z', 0, 0.0)",
rusqlite::params![mid, db_kind, text, normalized],
)
.expect("insert memory");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES (?1, ?2, ?3, 'project')",
rusqlite::params![mid, text, db_kind],
)
.expect("insert fts");
}
let weights = kimetsu_core::config::BrokerWeights::default();
let query = "cache convention failure".to_string();
let baseline = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: query.clone(),
budget_tokens: 4000,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("baseline retrieve");
let feature = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: query.clone(),
budget_tokens: 4000,
task_kind: TaskKind::Feature,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("feature retrieve");
let baseline_ids: Vec<&str> = baseline
.capsules
.iter()
.map(|c| c.expansion_handle.as_str())
.collect();
let feature_ids: Vec<&str> = feature
.capsules
.iter()
.map(|c| c.expansion_handle.as_str())
.collect();
assert_eq!(
baseline_ids, feature_ids,
"task_kind=Feature must produce identical retrieval to default; \
baseline={baseline_ids:?} feature={feature_ids:?}"
);
let baseline_scores: Vec<f32> = baseline.capsules.iter().map(|c| c.score).collect();
let feature_scores: Vec<f32> = feature.capsules.iter().map(|c| c.score).collect();
for (b, f) in baseline_scores.iter().zip(feature_scores.iter()) {
assert!(
(b - f).abs() < 1e-5,
"scores must be identical: baseline={b} feature={f}"
);
}
}
#[test]
fn debug_surfaces_more_failure_pattern_than_docs() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
for (i, text) in [
"auth token expired causes login failure",
"auth service crash on null pointer",
"auth regression after upgrade breaks sessions",
"auth error when certificate is invalid",
]
.iter()
.enumerate()
{
let mid = format!("mfp{i}");
let normalized = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score
)
VALUES (?1, 'project', 'failure_pattern', ?2, ?3, 1.0, NULL, '{}',
'2026-01-01T00:00:00Z', 0, 0.0)",
rusqlite::params![mid, text, normalized],
)
.expect("insert failure_pattern");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES (?1, ?2, 'failure_pattern', 'project')",
rusqlite::params![mid, text],
)
.expect("insert fts");
}
for (i, (db_kind, text)) in [
("convention", "auth module uses bearer tokens by convention"),
("convention", "auth scopes are documented in the API guide"),
("fact", "auth service runs on port 8443 in production"),
("fact", "auth uses JWT with RS256 signing for all tokens"),
]
.iter()
.enumerate()
{
let mid = format!("mconv{i}");
let normalized = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score
)
VALUES (?1, 'project', ?2, ?3, ?4, 1.0, NULL, '{}',
'2026-01-01T00:00:00Z', 0, 0.0)",
rusqlite::params![mid, db_kind, text, normalized],
)
.expect("insert convention/fact");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES (?1, ?2, ?3, 'project')",
rusqlite::params![mid, text, db_kind],
)
.expect("insert fts");
}
let weights = kimetsu_core::config::BrokerWeights::default();
let query = "auth token failure".to_string();
let debug_bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: query.clone(),
budget_tokens: 4000,
max_capsules: 4,
task_kind: TaskKind::Debug,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("debug retrieve");
let docs_bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&weights,
ContextRequest {
stage: "localization".to_string(),
query: query.clone(),
budget_tokens: 4000,
max_capsules: 4,
task_kind: TaskKind::Docs,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("docs retrieve");
let count_failure_pattern = |bundle: &ContextBundle| -> usize {
bundle
.capsules
.iter()
.filter(|c| capsule_matches_kind(c, "failure_pattern"))
.count()
};
let debug_fp = count_failure_pattern(&debug_bundle);
let docs_fp = count_failure_pattern(&docs_bundle);
assert!(
debug_fp > docs_fp,
"Debug must surface strictly more failure_pattern capsules than Docs: \
debug_fp={debug_fp} docs_fp={docs_fp}\n\
Debug capsules: {:?}\n\
Docs capsules: {:?}",
debug_bundle
.capsules
.iter()
.map(|c| format!("{}:{}", c.kind, &c.summary[..c.summary.len().min(60)]))
.collect::<Vec<_>>(),
docs_bundle
.capsules
.iter()
.map(|c| format!("{}:{}", c.kind, &c.summary[..c.summary.len().min(60)]))
.collect::<Vec<_>>(),
);
}
fn init_db_with_memory(memory_id: &str, text: &str) -> rusqlite::Connection {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let normalized = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score
)
VALUES (?1, 'project', 'fact', ?2, ?3, 1.0, NULL, '{}',
'2026-01-01T00:00:00Z', 0, 0.0)",
rusqlite::params![memory_id, text, normalized],
)
.expect("insert memory");
conn
}
#[test]
fn resolve_capsule_memory_returns_full_text() {
let conn = init_db_with_memory("test-mem-id", "Use rg over grep for speed");
let repo_root = std::path::Path::new("/fake-repo");
let result =
resolve_capsule(&conn, repo_root, "memory:test-mem-id").expect("should resolve");
assert_eq!(result, "Use rg over grep for speed");
}
#[test]
fn resolve_capsule_memory_missing_id_returns_err() {
let conn = init_db_with_memory("real-id", "some text");
let repo_root = std::path::Path::new("/fake-repo");
let err = resolve_capsule(&conn, repo_root, "memory:nonexistent-id")
.expect_err("should error for missing memory");
assert!(
err.to_string().contains("no active memory"),
"error message should mention missing: {err}"
);
}
#[test]
fn resolve_capsule_file_returns_bounded_content() {
let dir = make_test_dir("f2_file_resolve");
let content = "hello from the file\n";
std::fs::write(dir.join("notes.txt"), content).expect("write");
let result = resolve_capsule(
&rusqlite::Connection::open_in_memory().expect("open"),
&dir,
"file:notes.txt",
)
.expect("should resolve file");
assert!(result.contains("hello from the file"));
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn resolve_capsule_file_caps_large_file() {
let dir = make_test_dir("f2_file_cap");
let big = "A".repeat(FILE_EXPAND_CAP_BYTES * 3);
std::fs::write(dir.join("big.txt"), &big).expect("write");
let result = resolve_capsule(
&rusqlite::Connection::open_in_memory().expect("open"),
&dir,
"file:big.txt",
)
.expect("should resolve large file");
assert!(
result.len() <= FILE_EXPAND_CAP_BYTES + 200,
"result should be bounded: got {} bytes",
result.len()
);
assert!(
result.contains("truncated"),
"truncation marker should be present"
);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn resolve_capsule_unknown_handle_returns_err() {
let conn = rusqlite::Connection::open_in_memory().expect("open");
let err = resolve_capsule(&conn, std::path::Path::new("/r"), "blob:abc123")
.expect_err("should error");
assert!(
err.to_string().contains("unrecognised handle"),
"got: {err}"
);
}
#[test]
fn resolve_capsule_malformed_handle_returns_err() {
let conn = rusqlite::Connection::open_in_memory().expect("open");
let err = resolve_capsule(&conn, std::path::Path::new("/r"), "justnocolon")
.expect_err("should error");
assert!(
err.to_string().contains("unrecognised handle"),
"got: {err}"
);
}
#[test]
fn resolve_capsule_run_handle_returns_deferred_err() {
let conn = rusqlite::Connection::open_in_memory().expect("open");
let err = resolve_capsule(&conn, std::path::Path::new("/r"), "run:some-run-id")
.expect_err("run: should be deferred err");
assert!(err.to_string().contains("not yet supported"), "got: {err}");
}
#[test]
fn resolve_capsule_file_rejects_absolute_path() {
let conn = rusqlite::Connection::open_in_memory().expect("open");
let err = resolve_capsule(&conn, std::path::Path::new("/r"), "file:/etc/passwd")
.expect_err("should reject absolute path");
assert!(err.to_string().contains("absolute path"), "got: {err}");
}
fn make_capsule(summary: &str, score: f32) -> ContextCapsule {
ContextCapsule {
id: new_id().to_string(),
kind: "memory".to_string(),
summary: summary.to_string(),
token_estimate: 10,
expansion_handle: format!("memory:{}", new_id()),
provenance: vec![],
confidence: 1.0,
freshness: 1.0,
relevance: 1.0,
scope_weight: 1.0,
score,
superseded_hint: false,
rerank_policy_tier: 0,
claim_revision: None,
facts: vec![],
rerank_usefulness: None,
rerank_trust: None,
}
}
#[test]
fn rerank_capsules_reorders_by_query_overlap() {
use crate::embeddings::StubReranker;
let query = "rust async tokio";
let high_overlap = make_capsule("rust async tokio runtime", 0.0);
let low_overlap = make_capsule("python django framework", 0.0);
let capsules = vec![low_overlap.clone(), high_overlap.clone()];
let ranked = rerank_capsules(query, capsules, &StubReranker, 0.0, 0);
assert_eq!(ranked.len(), 2, "both capsules should survive (floor=0)");
assert!(
ranked[0].summary.contains("rust"),
"rust capsule must be first, got: {:?}",
ranked[0].summary
);
assert!(
ranked[0].score > 0.05,
"score must be overwritten by reranker: {}",
ranked[0].score
);
assert!(
ranked[0].score > ranked[1].score,
"high overlap must score higher: {} vs {}",
ranked[0].score,
ranked[1].score
);
}
#[test]
fn rerank_capsules_floor_drops_zero_overlap() {
use crate::embeddings::StubReranker;
let query = "rust async tokio";
let high = make_capsule("rust async tokio runtime", 0.0);
let zero = make_capsule("completely unrelated document xyz", 0.0);
let capsules = vec![high, zero];
let ranked = rerank_capsules(query, capsules, &StubReranker, 0.3, 0);
assert_eq!(ranked.len(), 1, "zero-overlap capsule must be dropped");
assert!(
ranked[0].summary.contains("rust"),
"only rust capsule should survive"
);
}
#[test]
fn rerank_capsules_cap_truncates() {
use crate::embeddings::StubReranker;
let query = "alpha beta gamma";
let capsules = vec![
make_capsule("alpha beta gamma delta", 0.0),
make_capsule("alpha beta", 0.0),
make_capsule("alpha", 0.0),
make_capsule("unrelated xyz", 0.0),
];
let ranked = rerank_capsules(query, capsules, &StubReranker, 0.0, 2);
assert_eq!(ranked.len(), 2, "cap=2 must truncate to 2 results");
assert!(
ranked[0].score >= ranked[1].score,
"results must be sorted descending"
);
}
#[test]
fn hardening_usefulness_cannot_dominate_relevance() {
let neutral = make_capsule("neutral", 0.0);
let mut useful = make_capsule("useful", 0.0);
useful.rerank_policy_tier = 1;
let out = rerank_capsules(
"q",
vec![neutral, useful],
&TwoScoreReranker(0.99, 0.31),
0.30,
0,
);
assert_eq!(out[0].summary, "neutral");
assert!(out[1].score <= 0.410001);
}
#[test]
fn hardening_rerank_preserves_decayed_usefulness_and_trust() {
let neutral = make_capsule("neutral", 0.0);
let mut stale_useful = make_capsule("stale", 0.0);
stale_useful.rerank_policy_tier = 1;
stale_useful.rerank_usefulness = Some(1.0001);
let out = rerank_capsules(
"q",
vec![neutral.clone(), stale_useful],
&TwoScoreReranker(0.9, 0.85),
0.0,
0,
);
assert_eq!(out[0].summary, "neutral");
let mut imported = make_capsule("imported", 0.0);
imported.rerank_usefulness = Some(1.5);
imported.rerank_trust = Some(0.5);
let out = rerank_capsules(
"q",
vec![neutral, imported],
&TwoScoreReranker(0.8, 0.9),
0.0,
0,
);
assert_eq!(out[0].summary, "neutral");
assert!((out[1].score - 0.5).abs() < 0.00001);
}
#[test]
fn hardening_freshness_has_thirty_day_half_life() {
let past = (OffsetDateTime::now_utc() - time::Duration::days(30))
.format(&time::format_description::well_known::Rfc3339)
.unwrap();
assert!((freshness(&past) - 0.5).abs() < 0.0001);
}
#[test]
fn rerank_reapplies_usefulness_but_not_to_superseded_capsules() {
let neutral = make_capsule("neutral", 0.0);
let mut useful = make_capsule("useful", 0.0);
useful.rerank_policy_tier = 1;
let out = rerank_capsules(
"q",
vec![neutral.clone(), useful.clone()],
&TwoScoreReranker(0.59, 0.50),
0.0,
0,
);
assert_eq!(out[0].summary, "useful", "usefulness survives reranking");
useful.superseded_hint = true;
useful.rerank_policy_tier = 1;
let out = rerank_capsules(
"q",
vec![neutral, useful],
&TwoScoreReranker(0.59, 0.50),
0.0,
0,
);
assert_eq!(
out[0].summary, "neutral",
"superseded memories must not keep their historic usefulness boost"
);
assert!(
(out[1].score - 0.40).abs() < 1e-6,
"supersession must be reapplied to the raw rerank score"
);
}
#[test]
fn rerank_capsules_fail_open_preserves_input_order() {
struct FailingReranker;
impl crate::embeddings::Reranker for FailingReranker {
fn rerank(
&self,
_query: &str,
_docs: &[&str],
) -> Result<Vec<f32>, crate::embeddings::EmbedderError> {
Err(crate::embeddings::EmbedderError::EmbedFailed(
"simulated failure".into(),
))
}
fn model_id(&self) -> &str {
"fail-reranker"
}
}
let query = "anything";
let c1 = make_capsule("first capsule", 0.9);
let c2 = make_capsule("second capsule", 0.5);
let c3 = make_capsule("third capsule", 0.1);
let capsules = vec![c1.clone(), c2.clone(), c3.clone()];
let out = rerank_capsules(query, capsules, &FailingReranker, 0.0, 0);
assert_eq!(out.len(), 3, "all capsules must be returned on error");
assert_eq!(out[0].summary, c1.summary, "order must be preserved");
assert_eq!(out[1].summary, c2.summary, "order must be preserved");
assert_eq!(out[2].summary, c3.summary, "order must be preserved");
}
#[test]
fn rerank_capsules_empty_input_returns_empty() {
use crate::embeddings::StubReranker;
let out = rerank_capsules("query", vec![], &StubReranker, 0.0, 0);
assert!(out.is_empty());
}
struct FixedReranker(f32);
impl crate::embeddings::Reranker for FixedReranker {
fn rerank(
&self,
_query: &str,
docs: &[&str],
) -> Result<Vec<f32>, crate::embeddings::EmbedderError> {
Ok(vec![self.0; docs.len()])
}
fn model_id(&self) -> &str {
"fixed-reranker"
}
}
struct TwoScoreReranker(f32, f32);
impl crate::embeddings::Reranker for TwoScoreReranker {
fn rerank(
&self,
_query: &str,
_docs: &[&str],
) -> Result<Vec<f32>, crate::embeddings::EmbedderError> {
Ok(vec![self.0, self.1])
}
fn model_id(&self) -> &str {
"two-score"
}
}
fn band_bundle(top_abs_evidence: f32) -> ContextBundle {
let mut capsule = make_capsule("a memory lesson", 0.9);
capsule.kind = "memory".to_string();
capsule.token_estimate = 10;
ContextBundle {
stage: "localization".into(),
budget_tokens: 4000,
used_tokens: 10,
capsules: vec![capsule],
excluded: vec![],
skipped: false,
top_score: 0.9,
top_abs_evidence,
evidence_coverage: 1.0,
uncovered_terms: vec![],
chronological: false,
known_fact_conflicts: vec![],
}
}
#[test]
fn band_arbitration_follows_the_cross_encoder() {
let approve = FixedReranker(ABSTAIN_RERANK_FLOOR + 0.2);
let out = rerank_and_arbitrate("q", band_bundle(0.50), Some(&approve), 0.55, 0.0, 0);
assert!(!out.skipped, "approved band bundle must inject");
assert_eq!(out.capsules.len(), 1);
let reject = FixedReranker(ABSTAIN_RERANK_FLOOR - 0.2);
let out = rerank_and_arbitrate("q", band_bundle(0.50), Some(&reject), 0.55, 0.0, 0);
assert!(out.skipped, "rejected band bundle must convert to skipped");
assert!(out.capsules.is_empty());
assert_eq!(out.used_tokens, 0);
assert_eq!(out.excluded.len(), 1, "rejected capsules land in excluded");
}
#[test]
fn band_arbitration_uses_raw_rerank_evidence() {
let mut bundle = band_bundle(0.50);
bundle.capsules[0].superseded_hint = true;
bundle
.capsules
.push(make_capsule("irrelevant distractor", 1.0));
let reranker = TwoScoreReranker(0.95, 0.0);
let out = rerank_and_arbitrate("q", bundle, Some(&reranker), 0.55, 0.0, 0);
assert!(!out.skipped, "raw rerank evidence above 0.9 must admit");
assert_eq!(out.capsules[0].summary, "a memory lesson");
assert!(
out.capsules[0].score < ABSTAIN_RERANK_FLOOR,
"the regression requires post-policy score below the raw-score floor"
);
}
#[test]
fn band_arbitration_uses_raw_evidence_before_policy_cap() {
let mut bundle = band_bundle(0.50);
bundle.capsules[0].rerank_policy_tier = 1;
bundle
.capsules
.push(make_capsule("high-confidence neutral", 1.0));
let reranker = TwoScoreReranker(0.50, 0.95);
let out = rerank_and_arbitrate("q", bundle, Some(&reranker), 0.55, 0.0, 1);
assert!(!out.skipped, "raw evidence outside the cap must admit");
assert_eq!(out.capsules.len(), 1);
assert_eq!(out.capsules[0].summary, "high-confidence neutral");
}
#[test]
fn band_arbitration_never_converts_out_of_band_bundles() {
let reject = FixedReranker(0.0);
let out = rerank_and_arbitrate("q", band_bundle(0.70), Some(&reject), 0.55, 0.0, 0);
assert!(
!out.skipped,
"evidence above the threshold is not arbitrated"
);
assert_eq!(out.capsules.len(), 1);
}
#[test]
fn band_fails_closed_without_a_reranker() {
let out = rerank_and_arbitrate("q", band_bundle(0.50), None, 0.55, 0.0, 0);
assert!(out.skipped, "band without an arbiter must abstain");
let out = rerank_and_arbitrate("q", band_bundle(0.70), None, 0.55, 0.0, 0);
assert!(!out.skipped);
let out = rerank_and_arbitrate("q", band_bundle(0.10), None, 0.0, 0.0, 0);
assert!(!out.skipped);
}
struct PositionReranker;
impl crate::embeddings::Reranker for PositionReranker {
fn rerank(
&self,
_query: &str,
docs: &[&str],
) -> Result<Vec<f32>, crate::embeddings::EmbedderError> {
Ok((0..docs.len()).map(|i| 0.99 - 0.01 * i as f32).collect())
}
fn model_id(&self) -> &str {
"position-reranker"
}
}
#[test]
fn supersession_penalty_survives_reranking() {
let mut old = make_capsule("deploy via make deploy-staging", 0.7);
old.superseded_hint = true; let new = make_capsule("deploy via make deploy-preview since the migration", 0.9);
let mut bundle = band_bundle(0.70); bundle.capsules = vec![old, new];
let out = rerank_and_arbitrate("q", bundle, Some(&PositionReranker), 0.55, 0.0, 0);
assert!(!out.skipped);
assert_eq!(out.capsules.len(), 2);
assert!(
out.capsules[0].summary.contains("deploy-preview"),
"the replacement must outrank the penalized incumbent after reranking; got {:?}",
out.capsules.iter().map(|c| &c.summary).collect::<Vec<_>>()
);
assert!(out.capsules[1].superseded_hint);
}
#[test]
fn band_spares_bundles_with_repo_evidence() {
let mut bundle = band_bundle(0.50);
let mut repo = make_capsule("README excerpt", 0.4);
repo.kind = "repo_file".to_string();
bundle.capsules.push(repo);
let reject = FixedReranker(0.0);
let out = rerank_and_arbitrate("q", bundle, Some(&reject), 0.55, 0.0, 0);
assert!(!out.skipped, "repo capsules suppress band conversion");
}
#[test]
fn compress_for_render_short_text_unchanged() {
let text = "project:fact - Use cargo fmt before committing.";
let out = compress_for_render(text, 3);
assert_eq!(out, text, "short text must not be altered");
}
#[test]
fn compress_for_render_strips_tags_prefix() {
let text = "[tags: rust, cargo] Always run cargo clippy before submitting a PR.";
let out = compress_for_render(text, 3);
assert!(
!out.starts_with('['),
"tags prefix must be stripped, got: {out:?}"
);
assert!(
out.contains("cargo clippy"),
"body must remain, got: {out:?}"
);
}
#[test]
fn compress_for_render_strips_context_suffix() {
let text =
"project:fact - Use cargo fmt. Always clippy clean. (context: Kimetsu brain lesson)";
let out = compress_for_render(text, 5);
assert!(
!out.contains("(context:"),
"context suffix must be stripped, got: {out:?}"
);
assert!(out.contains("cargo fmt"), "body must remain, got: {out:?}");
}
#[test]
fn compress_for_render_caps_sentences() {
let text =
"project:fact - First sentence. Second sentence. Third sentence. Fourth sentence.";
let out = compress_for_render(text, 2);
assert!(out.contains("First"), "first sentence must be present");
assert!(out.contains("Second"), "second sentence must be present");
assert!(
!out.contains("Third"),
"third sentence must be truncated, got: {out:?}"
);
}
#[test]
fn compress_for_render_preserves_scope_prefix() {
let text = "global_user:convention - First rule. Second rule. Third rule. Fourth rule.";
let out = compress_for_render(text, 2);
assert!(
out.starts_with("global_user:convention - "),
"scope prefix must be preserved, got: {out:?}"
);
assert!(out.contains("First"), "first sentence must remain");
assert!(!out.contains("Third"), "third sentence must be truncated");
}
#[test]
fn compress_for_render_empty_input_safe() {
let out = compress_for_render("", 3);
assert_eq!(out, "", "empty input must return empty string");
}
#[test]
fn compress_for_render_zero_max_sentences_returns_original() {
let text = "project:fact - Some lesson that is quite long. It keeps going. And going.";
let out = compress_for_render(text, 0);
assert_eq!(out, text);
}
#[test]
fn compress_for_render_utf8_safe() {
let text = "project:fact - こんにちは世界. Hello world. Third sentence. Fourth sentence.";
let out = compress_for_render(text, 2);
assert!(!out.is_empty(), "UTF-8 text must not produce empty output");
assert!(!out.contains("Third"), "third sentence must be truncated");
}
#[test]
fn compress_for_render_long_memory_reduces_tokens_by_25_percent() {
let long_summary = "project:fact - When a SQLite WAL file exists from a crashed process, \
opening the DB causes the WAL to be replayed. The replayed WAL may contain \
partial writes that corrupt the DB. Always check for WAL files before opening. \
Delete the WAL only after verifying the DB is consistent. Use PRAGMA integrity_check \
to validate after opening. If integrity_check fails, restore from backup. Never \
truncate the WAL without replaying it first. This pattern applies to any \
crash-recovery scenario.";
let raw_tokens = estimate_tokens(long_summary);
assert!(
raw_tokens > 60,
"test precondition: raw memory must be >60 tokens, got {raw_tokens}"
);
let compressed = compress_for_render(long_summary, 3);
let compressed_tokens = estimate_tokens(&compressed);
let reduction = 1.0 - (compressed_tokens as f64 / raw_tokens as f64);
assert!(
reduction >= 0.25,
"compression must reduce tokens by >=25% on long memories; \
raw={raw_tokens} compressed={compressed_tokens} reduction={reduction:.2}"
);
}
}
#[cfg(test)]
mod evidence_tests {
use super::*;
fn conn_with(texts: &[&str]) -> Connection {
let conn = Connection::open_in_memory().expect("open");
crate::schema::initialize(&conn).expect("schema");
for (i, text) in texts.iter().enumerate() {
conn.execute(
"INSERT INTO memories
(memory_id, scope, kind, text, normalized_text, confidence,
provenance_snapshot_json, created_at)
VALUES (?1, 'project', 'fact', ?2, ?2, 0.9, '{}', '2026-01-01T00:00:00Z')",
rusqlite::params![format!("m{i}"), text],
)
.expect("insert");
conn.execute("INSERT INTO memories_fts(memory_id,text,kind,scope) VALUES (?1,?2,'fact','project')",
params![format!("m{i}"),text]).unwrap();
}
conn
}
fn capsule(summary: &str) -> ContextCapsule {
ContextCapsule {
id: String::new(),
kind: "memory".to_string(),
summary: summary.to_string(),
token_estimate: 10,
expansion_handle: format!("memory:{summary}"),
provenance: Vec::new(),
confidence: 0.9,
freshness: 0.5,
relevance: 0.0,
scope_weight: 0.9,
score: 0.5,
superseded_hint: false,
rerank_policy_tier: 0,
claim_revision: None,
facts: vec![],
rerank_usefulness: None,
rerank_trust: None,
}
}
fn bundle(capsules: Vec<ContextCapsule>, coverage: f32, uncovered: &[&str]) -> ContextBundle {
ContextBundle {
stage: "localization".to_string(),
budget_tokens: 2000,
used_tokens: 20,
capsules,
excluded: Vec::new(),
skipped: false,
top_score: 0.7,
top_abs_evidence: 0.0,
evidence_coverage: coverage,
uncovered_terms: uncovered.iter().map(|s| s.to_string()).collect(),
chronological: false,
known_fact_conflicts: vec![],
}
}
#[test]
fn full_coverage_names_nothing() {
let conn = conn_with(&[
"checkpoint the wal before copying brain.db",
"vacuum reclaims dead pages",
]);
let (coverage, uncovered) = evidence_coverage(
&conn,
"checkpoint wal",
&[capsule(
"project:fact - checkpoint the wal before copying brain.db",
)],
);
assert!(coverage > 0.99, "got {coverage}");
assert!(uncovered.is_empty(), "got {uncovered:?}");
}
#[test]
fn partial_coverage_names_the_missing_terms() {
let conn = conn_with(&[
"checkpoint the wal before copying brain.db",
"the migration runner snapshots before each step",
]);
let (coverage, uncovered) = evidence_coverage(
&conn,
"checkpoint wal migration",
&[capsule(
"project:fact - checkpoint the wal before copying brain.db",
)],
);
assert!(coverage < 1.0, "coverage should be partial: {coverage}");
assert!(
uncovered.iter().any(|t| t.starts_with("migrat")),
"the uncovered term must be named: {uncovered:?}"
);
}
#[test]
fn coverage_is_collective_not_per_capsule() {
let conn = conn_with(&[
"checkpoint the wal before copying brain.db",
"the migration runner snapshots before each step",
]);
let (coverage, uncovered) = evidence_coverage(
&conn,
"checkpoint migration",
&[
capsule("project:fact - checkpoint the wal before copying"),
capsule("project:fact - the migration runner snapshots first"),
],
);
assert!(
coverage > 0.99,
"neither capsule covers both terms, but together they do: {coverage}"
);
assert!(uncovered.is_empty(), "got {uncovered:?}");
}
#[test]
fn an_unmeasurable_query_does_not_claim_a_gap() {
let conn = conn_with(&["checkpoint the wal"]);
let (coverage, uncovered) =
evidence_coverage(&conn, "the and of", &[capsule("project:fact - checkpoint")]);
assert_eq!(coverage, 1.0);
assert!(uncovered.is_empty());
}
#[test]
fn a_term_the_corpus_has_never_seen_counts_as_a_gap() {
let conn = conn_with(&[
"checkpoint the wal before copying brain.db",
"vacuum reclaims dead pages",
]);
let (coverage, uncovered) = evidence_coverage(
&conn,
"checkpoint the wal during a kubernetes rollout",
&[capsule(
"project:fact - checkpoint the wal before copying brain.db",
)],
);
assert!(
coverage <= PARTIAL_EVIDENCE_COVERAGE,
"an unknown half of the question must read as thin, not complete: {coverage}"
);
assert!(
uncovered.iter().any(|t| t.starts_with("kubernet")),
"the unknown term must be named: {uncovered:?}"
);
}
#[test]
fn a_ubiquitous_term_carries_no_weight() {
let conn = conn_with(&["kimetsu checkpoint wal", "kimetsu vacuum pages"]);
let (coverage, _) = evidence_coverage(
&conn,
"kimetsu vacuum",
&[capsule("project:fact - kimetsu vacuum pages")],
);
assert!(coverage > 0.99, "got {coverage}");
}
#[test]
fn an_empty_query_does_not_claim_a_gap() {
let conn = conn_with(&["checkpoint the wal"]);
assert_eq!(evidence_coverage(&conn, "", &[]).0, 1.0);
}
#[test]
fn a_complete_bundle_gets_no_notice() {
assert!(partial_evidence_notice(&bundle(vec![capsule("a")], 1.0, &[])).is_none());
assert!(
partial_evidence_notice(&bundle(vec![capsule("a")], 0.9, &["x"])).is_none(),
"above the threshold is not partial"
);
}
#[test]
fn an_empty_or_skipped_bundle_gets_no_notice() {
let mut skipped = bundle(Vec::new(), 0.0, &["x"]);
skipped.skipped = true;
assert!(
partial_evidence_notice(&skipped).is_none(),
"an empty bundle already says everything it can"
);
assert!(partial_evidence_notice(&bundle(Vec::new(), 0.0, &["x"])).is_none());
}
#[test]
fn a_partial_bundle_names_what_is_missing_and_tells_the_reader_what_to_do() {
let notice =
partial_evidence_notice(&bundle(vec![capsule("a")], 0.3, &["migration", "rollback"]))
.expect("a thin bundle must be flagged");
assert!(notice.contains("migration"), "got: {notice}");
assert!(notice.contains("rollback"), "got: {notice}");
assert!(
notice.contains("unknown"),
"the notice must tell the reader to abstain, not just report a gap: {notice}"
);
}
#[test]
fn the_notice_caps_how_many_terms_it_names() {
let terms: Vec<String> = (0..12).map(|i| format!("term{i}")).collect();
let refs: Vec<&str> = terms.iter().map(String::as_str).collect();
let notice =
partial_evidence_notice(&bundle(vec![capsule("a")], 0.1, &refs)).expect("flagged");
assert!(notice.contains("and 6 more"), "got: {notice}");
assert!(!notice.contains("term9"), "got: {notice}");
}
#[test]
fn the_y_ies_pair_shares_a_stem() {
for (a, b) in [
("retry", "retries"),
("query", "queries"),
("policy", "policies"),
("memory", "memories"),
("binary", "binaries"),
("registry", "registries"),
] {
assert_eq!(
light_stem(a),
light_stem(b),
"{a}/{b} stemmed to {:?}/{:?}",
light_stem(a),
light_stem(b)
);
}
}
#[test]
fn a_vowel_y_is_not_stripped() {
assert_eq!(light_stem("delay"), "delay");
assert_eq!(light_stem("gateway"), "gateway");
assert_eq!(light_stem("journeys"), "journey");
}
#[test]
fn short_words_keep_their_ending() {
assert_eq!(light_stem("body"), "body");
assert_eq!(light_stem("copy"), "copy");
}
#[test]
fn the_original_suffix_rules_still_hold() {
assert_eq!(light_stem("benchmarked"), "benchmark");
assert_eq!(light_stem("benchmarking"), "benchmark");
assert_eq!(light_stem("migrations"), "migration");
assert_eq!(light_stem("run"), "run");
}
#[test]
fn an_inflected_corpus_term_counts_as_covered() {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
let text = "the ingest worker retries a failed batch three times before giving up";
let normalized = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"
INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at
)
VALUES ('m_retry', 'project', 'fact', ?1, ?2, 1.0, NULL, '{}',
'2026-01-01T00:00:00Z')
",
rusqlite::params![text, normalized],
)
.expect("insert memory");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES ('m_retry', ?1, 'fact', 'project')",
rusqlite::params![text],
)
.expect("insert fts");
let bundle = retrieve_context_with_embedder(
&conn,
"/fake-repo",
&kimetsu_core::config::BrokerWeights::default(),
ContextRequest {
stage: "localization".to_string(),
query: "how many times does the ingest worker retry a failed batch".to_string(),
budget_tokens: 4000,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve");
assert!(
!bundle.uncovered_terms.iter().any(|t| t.starts_with("retr")),
"`retry` must match a corpus that says `retries`; uncovered: {:?}",
bundle.uncovered_terms
);
}
fn ordering_conn() -> rusqlite::Connection {
let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
crate::schema::initialize(&conn).expect("init schema");
for (mid, created, text) in [
(
"m_late",
"2026-06-01T09:00:00Z",
"switched the error type to thiserror",
),
(
"m_early",
"2026-01-15T10:00:00Z",
"ran the thiserror schema migration",
),
] {
let normalized = kimetsu_core::memory::normalize_memory_text(text);
conn.execute(
"
INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at
)
VALUES (?1, 'project', 'fact', ?2, ?3, 1.0, NULL, '{}', ?4)
",
rusqlite::params![mid, text, normalized, created],
)
.expect("insert memory");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES (?1, ?2, 'fact', 'project')",
rusqlite::params![mid, text],
)
.expect("insert fts");
}
conn
}
fn ordering_bundle(conn: &rusqlite::Connection, query: &str) -> ContextBundle {
retrieve_context_with_embedder(
conn,
"/fake-repo",
&kimetsu_core::config::BrokerWeights::default(),
ContextRequest {
stage: "localization".to_string(),
query: query.to_string(),
budget_tokens: 4000,
..Default::default()
},
&[],
&embeddings::NoopEmbedder,
)
.expect("retrieve")
}
#[test]
fn an_ordering_query_returns_a_dated_chronological_bundle() {
let conn = ordering_conn();
let bundle = ordering_bundle(&conn, "did we run the thiserror migration before or after");
assert!(bundle.chronological, "the query asked about order");
let order: Vec<&str> = bundle
.capsules
.iter()
.filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
.collect();
assert_eq!(order, vec!["m_early", "m_late"], "oldest first");
for (capsule, date) in bundle.capsules.iter().zip(["2026-01-15", "2026-06-01"]) {
assert!(
capsule.summary.contains(&format!("[{date}]")),
"every capsule carries its date; got: {}",
capsule.summary
);
}
}
#[test]
fn an_ordinary_query_is_untouched() {
let conn = ordering_conn();
let bundle = ordering_bundle(&conn, "how do we handle thiserror errors");
assert!(!bundle.chronological);
for capsule in &bundle.capsules {
assert!(
!capsule.summary.contains('['),
"no dates on a non-ordering query; got: {}",
capsule.summary
);
}
}
#[test]
fn ordering_changes_the_rendering_not_the_selection() {
let conn = ordering_conn();
let ordered = ordering_bundle(&conn, "did we run the thiserror migration before or after");
let plain = ordering_bundle(&conn, "did we run the thiserror migration");
let mut got: Vec<&str> = ordered
.capsules
.iter()
.map(|c| c.expansion_handle.as_str())
.collect();
let mut want: Vec<&str> = plain
.capsules
.iter()
.map(|c| c.expansion_handle.as_str())
.collect();
got.sort_unstable();
want.sort_unstable();
assert_eq!(got, want, "same capsules, different order");
}
#[test]
fn the_dates_are_counted_against_the_budget() {
let conn = ordering_conn();
let ordered = ordering_bundle(&conn, "did we run the thiserror migration before or after");
let plain = ordering_bundle(&conn, "did we run the thiserror migration");
assert!(
ordered.used_tokens > plain.used_tokens,
"dated: {} vs plain: {}",
ordered.used_tokens,
plain.used_tokens
);
assert_eq!(
ordered.used_tokens,
ordered
.capsules
.iter()
.map(|c| c.token_estimate)
.sum::<u32>(),
"used_tokens must match what was actually rendered"
);
}
}
#[cfg(test)]
mod hardening_tests {
use super::*;
fn corpus() -> Connection {
let conn = Connection::open_in_memory().unwrap();
crate::schema::initialize(&conn).unwrap();
for (id, text) in [
("live", "routing routing routes"),
("future", "routing"),
("expired", "routing"),
("offset", "routing"),
("other", "rerouting unrelated"),
] {
conn.execute("INSERT INTO memories (memory_id,scope,kind,text,normalized_text,confidence,created_at,provenance_snapshot_json)
VALUES (?1,'project','fact',?2,?2,1,'2020-01-01T00:00:00Z','{}')", params![id,text]).unwrap();
conn.execute("INSERT INTO memories_fts(memory_id,text,kind,scope) VALUES (?1,?2,'fact','project')",params![id,text]).unwrap();
}
conn
}
#[test]
fn hardening_live_lexical_and_recency_apply_both_time_bounds() {
let conn = corpus();
let now = OffsetDateTime::now_utc();
let fmt = &time::format_description::well_known::Rfc3339;
let future = (now + time::Duration::hours(1)).format(fmt).unwrap();
let expired = (now - time::Duration::seconds(2)).format(fmt).unwrap();
let offset = (now - time::Duration::seconds(2))
.to_offset(time::UtcOffset::from_hms(12, 0, 0).unwrap())
.format(fmt)
.unwrap();
conn.execute(
"UPDATE memories SET valid_from=?1 WHERE memory_id='future'",
params![future],
)
.unwrap();
conn.execute(
"UPDATE memories SET valid_to=?1 WHERE memory_id='expired'",
params![expired],
)
.unwrap();
conn.execute(
"UPDATE memories SET valid_to=?1 WHERE memory_id='offset'",
params![offset],
)
.unwrap();
for candidates in [
memory_fts_candidates(
&conn,
&["routing".into()],
"routing*",
80,
None,
30.0,
false,
)
.unwrap(),
latest_memory_candidates(&conn, &["routing".into()], 200, None, 30.0, false).unwrap(),
] {
let ids: Vec<_> = candidates
.iter()
.map(|c| c.capsule.expansion_handle.as_str())
.collect();
assert!(ids.contains(&"memory:live"));
for id in ["memory:future", "memory:expired", "memory:offset"] {
assert!(!ids.contains(&id), "returned {id}");
}
}
}
#[test]
fn hardening_hydration_binds_text_revision_before_later_correction() {
let conn = corpus();
let candidates = memory_fts_candidates(
&conn,
&["routing".into()],
"routing*",
80,
None,
30.0,
false,
)
.unwrap();
let capsules: Vec<_> = candidates.into_iter().map(|c| c.capsule).collect();
assert_eq!(memory_revision_bindings(&capsules)["live"], "baseline:live");
conn.execute("INSERT INTO memory_revisions(memory_id,event_id,text,kind,known_at,effective_at,confidence,use_count,usefulness_score)
VALUES ('live','corrected','changed claim','fact','2026-01-01T00:00:00Z','2026-01-01T00:00:00Z',1,0,0)",[]).unwrap();
conn.execute(
"UPDATE memories SET text='changed claim' WHERE memory_id='live'",
[],
)
.unwrap();
assert_eq!(
crate::projector::claim_revision_at(&conn, "live", None).unwrap(),
"corrected"
);
assert_eq!(memory_revision_bindings(&capsules)["live"], "baseline:live");
assert!(
capsules
.iter()
.find(|c| c.expansion_handle == "memory:live")
.unwrap()
.summary
.contains("routing routing")
);
}
#[cfg(feature = "embeddings")]
#[test]
fn hardening_ann_hydration_filters_time_bounds() {
let conn = corpus();
let blob = crate::embeddings::encode_embedding(&[1.0, 0.0]);
conn.execute(
"UPDATE memories SET embedding=?1,embedding_model='test'",
params![blob],
)
.unwrap();
conn.execute(
"UPDATE memories SET valid_from='2099-01-01T00:00:00Z' WHERE memory_id='future'",
[],
)
.unwrap();
let expired = (OffsetDateTime::now_utc() - time::Duration::seconds(2))
.to_offset(time::UtcOffset::from_hms(12, 0, 0).unwrap())
.format(&time::format_description::well_known::Rfc3339)
.unwrap();
conn.execute(
"UPDATE memories SET valid_to=?1 WHERE memory_id IN ('expired','offset')",
params![expired],
)
.unwrap();
let qe = QueryEmbedding {
vector: vec![1.0, 0.0],
model_id: "test".into(),
};
let out = memory_ann_candidates(&conn, &qe, 80, &["routing".into()], 30.0, false).unwrap();
assert_eq!(out.len(), 2);
for c in out {
assert!(matches!(
c.capsule.expansion_handle.as_str(),
"memory:live" | "memory:other"
));
}
}
#[test]
fn hardening_idf_counts_prefix_documents_not_occurrences_or_substrings() {
let conn = corpus();
let tokens = vec!["rout".into(), "absent".into()];
let coverage = coverage_token_idf(&conn, &tokens).unwrap();
assert!((coverage["rout"] - (6.0_f32 / 5.0).ln()).abs() < 0.00001);
assert!((coverage["absent"] - 6.0_f32.ln()).abs() < 0.00001);
assert_eq!(corpus_token_idf(&conn, &tokens).unwrap()["absent"], 0.0);
}
}
#[cfg(test)]
mod structured_fact_hydration_tests {
use super::*;
use kimetsu_core::{event::Event, ids::RunId};
#[test]
fn lexical_and_recency_capsules_keep_their_delivered_fact_revision() {
let c = Connection::open_in_memory().unwrap();
crate::schema::initialize(&c).unwrap();
crate::projector::apply_events(&c,&[Event::new(RunId::new(),"memory.accepted",serde_json::json!({
"memory_id":"m","scope":"project","kind":"fact","text":"Orchid staging gateway port is 7319."
}))]).unwrap();
let mut delivered = Vec::new();
for candidates in [
memory_fts_candidates(&c, &["orchid".into()], "orchid*", 80, None, 30.0, true).unwrap(),
latest_memory_candidates(&c, &["orchid".into()], 200, None, 30.0, true).unwrap(),
] {
let capsule = &candidates[0].capsule;
assert_eq!(capsule.facts.len(), 1);
assert_eq!(capsule.facts[0].claim.value, "7319");
assert_eq!(
capsule.claim_revision.as_deref(),
Some(capsule.facts[0].claim_revision.as_str())
);
delivered.push(capsule.clone());
}
crate::projector::apply_events(
&c,
&[Event::new(
RunId::new(),
"memory.corrected",
serde_json::json!({
"memory_id":"m","text":"Orchid staging gateway port is 8420."
}),
)],
)
.unwrap();
for capsule in delivered {
assert!(capsule.summary.contains("7319"));
assert_eq!(capsule.facts[0].claim.value, "7319");
}
let latest =
latest_memory_candidates(&c, &["orchid".into()], 200, None, 30.0, true).unwrap();
assert_eq!(latest[0].capsule.facts[0].claim.value, "8420");
}
#[test]
fn legacy_wire_capsules_default_to_empty_fact_evidence() {
let c = ContextCapsule::wire_minimal("hello".into(), "memory".into(), 1.0);
let json = serde_json::to_value(&c).unwrap();
assert!(json.get("facts").is_none());
assert!(
serde_json::from_value::<ContextCapsule>(json)
.unwrap()
.facts
.is_empty()
);
}
}
#[cfg(test)]
mod disabled_fact_hydration_tests {
use super::*;
#[test]
fn ordinary_retrieval_does_not_read_the_fact_projection() {
let c = Connection::open_in_memory().unwrap();
crate::schema::initialize(&c).unwrap();
crate::projector::apply_events(
&c,
&[kimetsu_core::event::Event::new(
kimetsu_core::ids::RunId::new(),
"memory.accepted",
serde_json::json!({"memory_id":"m","text":"Orchid gateway port is 7319."}),
)],
)
.unwrap();
c.execute_batch("DROP TABLE memory_facts").unwrap();
for query in ["Orchid", ""] {
let out =
memory_candidates_flat(&c, query, None, 30.0, crate::fusion::Fusion::Linear, false)
.unwrap();
assert_eq!(out.len(), 1);
assert!(out[0].capsule.facts.is_empty());
}
}
}
#[cfg(test)]
mod deferred_fact_budget_tests {
use super::*;
#[test]
fn initial_retrieval_budget_must_not_hide_an_eligible_conflicting_fact() {
let c = Connection::open_in_memory().unwrap();
crate::schema::initialize(&c).unwrap();
for (id, value) in [("a", "7319"), ("b", "7320")] {
let text = format!(
"Orchid gateway port is {value}. Stable operation. Recorded settings. {}",
"Operational notes remain available. ".repeat(350)
);
crate::projector::apply_events(
&c,
&[kimetsu_core::event::Event::new(
kimetsu_core::ids::RunId::new(),
"memory.accepted",
serde_json::json!({"memory_id":id,"scope":"project","kind":"fact","text":text}),
)],
)
.unwrap();
}
let query = "What is the Orchid gateway port?";
let policy = crate::serving::ServingPolicy {
budget: 6000,
cap: 1,
explicit_fact_guard: true,
..Default::default()
};
let request = ContextRequest {
stage: "localization".into(),
query: query.into(),
budget_tokens: 6000,
..Default::default()
};
let weights = BrokerWeights::default();
let mut ordinary = request.clone();
ordinary.max_capsules = 6;
let ordinary = retrieve_context_with_embedder(
&c,
"/fake-repo",
&weights,
ordinary,
&[],
&crate::embeddings::NoopEmbedder,
)
.unwrap();
assert_eq!(ordinary.capsules.len(), 1);
assert!(ordinary.used_tokens <= 3000);
let selected = retrieve_context_with_embedder(
&c,
"/fake-repo",
&weights,
policy.prepare(request, false),
&[],
&crate::embeddings::NoopEmbedder,
)
.unwrap();
assert_eq!(
selected.capsules.len(),
2,
"both eligible claims must reach arbitration before delivery budgeting"
);
let selected = policy.arbitrate(query, selected, None, 0.0);
let delivered =
policy.render_for_query(query, selected, true, crate::serving::EVAL_EXPOSURE_ID);
assert_eq!(delivered.capsules.len(), 1);
assert_eq!(delivered.payload["answerability"]["status"], "conflicting");
assert!(delivered.payload["used_tokens"].as_u64().unwrap() <= 6000);
}
}