1use std::cmp::Ordering;
2use std::collections::HashMap;
3
4use kimetsu_core::config::{BrokerWeights, StageWeights};
5use kimetsu_core::memory::MemoryScope;
6use kimetsu_core::{KimetsuResult, ids::new_id};
7use rusqlite::{Connection, OptionalExtension, params};
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
25pub enum TaskKind {
26 #[default]
29 Feature,
30 Debug,
33 Refactor,
36 Docs,
39 Investigation,
42}
43
44pub fn classify_task(task: &str) -> TaskKind {
50 let lower = task.to_ascii_lowercase();
51
52 const DEBUG_KW: &[&str] = &[
54 "fix",
55 "bug",
56 "error",
57 "fail",
58 "crash",
59 "panic",
60 "regression",
61 "broken",
62 "debug",
63 "stack trace",
64 "exception",
65 ];
66 if DEBUG_KW.iter().any(|kw| lower.contains(kw)) {
67 return TaskKind::Debug;
68 }
69
70 const INVESTIGATE_KW: &[&str] = &[
72 "investigate",
73 "analyze",
74 "understand",
75 " why ",
76 "explore",
77 "find out",
78 "root cause",
79 "audit",
80 "trace",
81 ];
82 if INVESTIGATE_KW.iter().any(|kw| lower.contains(kw)) {
83 return TaskKind::Investigation;
84 }
85
86 const REFACTOR_KW: &[&str] = &[
88 "refactor",
89 "rename",
90 "cleanup",
91 "clean up",
92 "restructure",
93 "simplify",
94 "extract",
95 "deduplicate",
96 "reorganize",
97 ];
98 if REFACTOR_KW.iter().any(|kw| lower.contains(kw)) {
99 return TaskKind::Refactor;
100 }
101
102 const DOCS_KW: &[&str] = &[
104 "document",
105 "readme",
106 "changelog",
107 "comment",
108 "docstring",
109 "docs",
110 "tutorial",
111 "guide",
112 ];
113 if DOCS_KW.iter().any(|kw| lower.contains(kw)) {
114 return TaskKind::Docs;
115 }
116
117 TaskKind::Feature
119}
120
121fn weights_for_task_kind(base: StageWeights, kind: TaskKind) -> StageWeights {
137 match kind {
138 TaskKind::Feature => base,
139 TaskKind::Debug => renorm(StageWeights {
140 freshness: base.freshness * 1.6,
141 ..base
142 }),
143 TaskKind::Refactor => renorm(StageWeights {
144 scope: base.scope * 1.6,
145 ..base
146 }),
147 TaskKind::Investigation => renorm(StageWeights {
148 relevance: base.relevance * 1.4,
149 ..base
150 }),
151 TaskKind::Docs => renorm(StageWeights {
152 confidence: base.confidence * 1.15,
153 ..base
154 }),
155 }
156}
157
158fn renorm(w: StageWeights) -> StageWeights {
162 let sum = w.relevance + w.confidence + w.freshness + w.scope;
163 if sum <= f32::EPSILON {
164 return w;
165 }
166 StageWeights {
175 relevance: w.relevance / sum,
176 confidence: w.confidence / sum,
177 freshness: w.freshness / sum,
178 scope: w.scope / sum,
179 }
180}
181
182fn task_kind_prefer_roles(kind: TaskKind) -> &'static [&'static str] {
188 match kind {
189 TaskKind::Feature => &[],
190 TaskKind::Debug => &["failure_pattern"],
191 TaskKind::Refactor => &["convention"],
192 TaskKind::Investigation => &["fact", "preference"],
193 TaskKind::Docs => &["convention"],
194 }
195}
196use time::OffsetDateTime;
197
198use crate::embeddings::{
199 self, DEFAULT_HYBRID_ALPHA, Embedder, cosine_similarity, decode_embedding,
200};
201
202#[derive(Debug, Clone)]
210pub(crate) struct QueryEmbedding {
211 pub(crate) vector: Vec<f32>,
212 pub(crate) model_id: String,
213}
214
215impl QueryEmbedding {
216 fn from_embedder(embedder: &dyn Embedder, query: &str) -> Option<Self> {
217 if embedder.is_noop() {
218 return None;
219 }
220 match embedder.embed(query) {
221 Ok(v) if v.len() == embedder.dim() => Some(Self {
222 vector: v,
223 model_id: embedder.model_id().to_string(),
224 }),
225 _ => None,
230 }
231 }
232}
233
234#[derive(Debug, Clone, Serialize, Deserialize)]
235pub struct ContextCapsule {
236 pub id: String,
237 pub kind: String,
238 pub summary: String,
239 pub token_estimate: u32,
240 pub expansion_handle: String,
241 pub provenance: Vec<ProvenanceRef>,
242 pub confidence: f32,
243 pub freshness: f32,
244 pub relevance: f32,
245 pub scope_weight: f32,
246 pub score: f32,
247}
248
249impl ContextCapsule {
250 pub fn wire_minimal(summary: String, kind: String, score: f32) -> Self {
254 Self {
255 id: String::new(),
256 kind,
257 summary,
258 token_estimate: 0,
259 expansion_handle: String::new(),
260 provenance: Vec::new(),
261 confidence: 0.0,
262 freshness: 0.0,
263 relevance: 0.0,
264 scope_weight: 0.0,
265 score,
266 }
267 }
268}
269
270#[derive(Debug, Clone, Serialize, Deserialize)]
271pub struct ProvenanceRef {
272 pub source: String,
273 pub id: String,
274 pub excerpt: Option<String>,
275}
276
277#[derive(Debug, Clone, Default)]
278pub struct ContextRequest {
279 pub stage: String,
280 pub query: String,
281 pub budget_tokens: u32,
282 pub fusion: String,
290 pub normalization: String,
298 pub tags: Vec<String>,
303 pub min_score: f32,
308 pub max_capsules: usize,
311 pub prefer_roles: Vec<String>,
315 pub kinds: Vec<String>,
322 pub min_semantic_score: f32,
330 pub min_lexical_coverage: f32,
340 pub task_kind: TaskKind,
345}
346
347#[derive(Debug, Clone)]
348pub struct ContextBundle {
349 pub stage: String,
350 pub budget_tokens: u32,
351 pub used_tokens: u32,
352 pub capsules: Vec<ContextCapsule>,
353 pub excluded: Vec<ContextCapsule>,
354 pub skipped: bool,
357 pub top_score: f32,
360 pub evidence_coverage: f32,
375 pub uncovered_terms: Vec<String>,
381 pub chronological: bool,
388}
389
390fn coverage_token_idf(conn: &Connection, tokens: &[String]) -> KimetsuResult<HashMap<String, f32>> {
407 let mut idf = HashMap::new();
408 let n: i64 = conn
409 .query_row(
410 "SELECT COUNT(*) FROM memories WHERE invalidated_at IS NULL",
411 [],
412 |row| row.get(0),
413 )
414 .unwrap_or(0);
415 if n == 0 {
416 return Ok(idf);
417 }
418 let mut stmt = conn.prepare_cached(
419 "SELECT COUNT(*) FROM memories \
420 WHERE invalidated_at IS NULL AND lower(text) LIKE ?1 ESCAPE '\\'",
421 )?;
422 for token in tokens {
423 let pattern = format!("%{}%", escape_like(token));
424 let df: i64 = stmt
425 .query_row(params![pattern], |row| row.get(0))
426 .unwrap_or(0);
427 let weight = (((n + 1) as f32) / ((df + 1) as f32)).ln().max(0.0);
429 idf.insert(token.clone(), weight);
430 }
431 Ok(idf)
432}
433
434pub fn partial_evidence_notice(bundle: &ContextBundle) -> Option<String> {
443 if bundle.skipped || bundle.capsules.is_empty() {
444 return None; }
446 if bundle.evidence_coverage > PARTIAL_EVIDENCE_COVERAGE || bundle.uncovered_terms.is_empty() {
447 return None;
448 }
449 const MAX_NAMED: usize = 6;
453 let named: Vec<&str> = bundle
454 .uncovered_terms
455 .iter()
456 .take(MAX_NAMED)
457 .map(String::as_str)
458 .collect();
459 let more = bundle.uncovered_terms.len().saturating_sub(named.len());
460 let suffix = if more > 0 {
461 format!(" (and {more} more)")
462 } else {
463 String::new()
464 };
465 Some(format!(
466 "Partial memory: nothing above covers {}{}. Treat the rest as unknown \
467 rather than inferring it.",
468 named.join(", "),
469 suffix
470 ))
471}
472
473pub const PARTIAL_EVIDENCE_COVERAGE: f32 = 0.5;
479
480pub(crate) fn evidence_coverage(
487 conn: &Connection,
488 query: &str,
489 capsules: &[ContextCapsule],
490) -> (f32, Vec<String>) {
491 let content = content_tokens(query);
492 if content.is_empty() {
493 return (1.0, Vec::new());
494 }
495 let Ok(idf) = coverage_token_idf(conn, &content) else {
496 return (1.0, Vec::new());
497 };
498 let haystack = capsules
500 .iter()
501 .map(|c| c.summary.to_ascii_lowercase())
502 .collect::<Vec<_>>()
503 .join(" ");
504
505 let mut total = 0.0f32;
506 let mut hit = 0.0f32;
507 let mut uncovered = Vec::new();
508 for token in &content {
509 let weight = idf.get(token).copied().unwrap_or(0.0);
510 if weight <= 0.0 {
511 continue; }
513 total += weight;
514 if haystack.contains(token.as_str()) {
515 hit += weight;
516 } else {
517 uncovered.push(token.clone());
518 }
519 }
520 if total <= f32::EPSILON {
521 return (1.0, Vec::new());
524 }
525 (hit / total, uncovered)
526}
527
528#[derive(Debug, Clone)]
534pub(crate) struct Candidate {
535 pub(crate) capsule: ContextCapsule,
536 pub(crate) raw_relevance: f32,
537 pub(crate) embedding: Option<Vec<f32>>,
543 pub(crate) cosine: Option<f32>,
547 pub(crate) created_at: Option<String>,
552}
553
554pub fn retrieve_context(
555 conn: &Connection,
556 repo_root: &str,
557 weights: &BrokerWeights,
558 request: ContextRequest,
559) -> KimetsuResult<ContextBundle> {
560 retrieve_context_multi(conn, repo_root, weights, request, &[])
561}
562
563pub fn retrieve_context_multi(
580 conn: &Connection,
581 repo_root: &str,
582 weights: &BrokerWeights,
583 request: ContextRequest,
584 extra_memory_conns: &[&Connection],
585) -> KimetsuResult<ContextBundle> {
586 let embedder = embeddings::open_default_embedder();
587 retrieve_context_with_embedder(
588 conn,
589 repo_root,
590 weights,
591 request,
592 extra_memory_conns,
593 embedder,
594 )
595}
596
597pub fn retrieve_context_with_embedder(
609 conn: &Connection,
610 repo_root: &str,
611 weights: &BrokerWeights,
612 request: ContextRequest,
613 extra_memory_conns: &[&Connection],
614 embedder: &dyn Embedder,
615) -> KimetsuResult<ContextBundle> {
616 retrieve_context_with_embedder_and_backend(
617 conn,
618 repo_root,
619 weights,
620 request,
621 extra_memory_conns,
622 embedder,
623 &crate::backend::FlatBackend {
624 fusion: crate::fusion::Fusion::Linear,
625 },
626 )
627}
628
629pub(crate) fn retrieve_context_with_embedder_and_backend(
642 conn: &Connection,
643 repo_root: &str,
644 weights: &BrokerWeights,
645 request: ContextRequest,
646 extra_memory_conns: &[&Connection],
647 embedder: &dyn Embedder,
648 backend: &dyn crate::backend::RetrievalBackend,
649) -> KimetsuResult<ContextBundle> {
650 let query_embedding = QueryEmbedding::from_embedder(embedder, &request.query);
651 let half_life_days = weights.decay_half_life_days;
652 let mut candidates = Vec::new();
653 candidates.extend(backend.memory_candidates(
654 conn,
655 &request.query,
656 query_embedding.as_ref(),
657 half_life_days,
658 )?);
659 for extra in extra_memory_conns {
660 candidates.extend(backend.memory_candidates(
661 extra,
662 &request.query,
663 query_embedding.as_ref(),
664 half_life_days,
665 )?);
666 }
667 crate::reinforce::apply_query_routing(
673 conn,
674 &request.query,
675 query_embedding.as_ref(),
676 &mut candidates,
677 );
678
679 candidates.extend(repo_file_candidates(conn, repo_root, &request.query, 30)?);
680 candidates.extend(manifest_candidates(conn, repo_root, &request.query)?);
681
682 if !request.kinds.is_empty() {
689 candidates.retain(|c| {
690 request
691 .kinds
692 .iter()
693 .any(|k| capsule_matches_kind(&c.capsule, k))
694 });
695 }
696
697 if request.min_lexical_coverage > 0.0 {
714 let content = content_tokens(&request.query);
715 if !content.is_empty() {
716 let idf = corpus_token_idf(conn, &content)?;
717 let total_idf: f32 = content
718 .iter()
719 .map(|t| idf.get(t).copied().unwrap_or(0.0))
720 .sum();
721 if total_idf > f32::EPSILON {
724 candidates.retain(|c| {
725 if c.capsule.kind != "memory" {
726 return true; }
728 if c.cosine.is_some_and(|cos| cos >= SEMANTIC_KEEP_COSINE) {
731 return true;
732 }
733 weighted_coverage(&content, &idf, &c.capsule.summary)
734 >= request.min_lexical_coverage
735 });
736 }
737 }
738 }
739
740 let stage_weights = weights_for_stage(weights, &request.stage);
743 let effective_weights = weights_for_task_kind(stage_weights, request.task_kind);
744 normalize_and_score(
745 &mut candidates,
746 effective_weights,
747 Normalization::from_config(&request.normalization),
748 );
749
750 let kind_role_hints = task_kind_prefer_roles(request.task_kind);
753 let mut effective_prefer_roles: Vec<String> = request.prefer_roles.clone();
754 for &hint in kind_role_hints {
755 let hint_s = hint.to_string();
756 if !effective_prefer_roles.contains(&hint_s) {
757 effective_prefer_roles.push(hint_s);
758 }
759 }
760
761 if !request.tags.is_empty() || !effective_prefer_roles.is_empty() {
773 let tags_lc: Vec<String> = request
774 .tags
775 .iter()
776 .map(|t| t.to_ascii_lowercase())
777 .collect();
778 for c in &mut candidates {
779 let summary_lc = c.capsule.summary.to_ascii_lowercase();
780 if !tags_lc.is_empty() && tags_lc.iter().any(|t| summary_lc.contains(t.as_str())) {
781 c.capsule.score *= 1.4;
782 }
783 if !effective_prefer_roles.is_empty()
784 && effective_prefer_roles.iter().any(|r| {
785 if c.capsule.kind == "memory" {
793 capsule_matches_kind(&c.capsule, r.as_str())
794 } else {
795 c.capsule.kind.contains(r.as_str())
796 }
797 })
798 {
799 c.capsule.score *= 1.3;
800 }
801 }
802 }
803
804 if query_embedding.is_some() && request.min_semantic_score > 0.0 {
819 candidates.retain(|c| {
820 match c.cosine {
823 Some(cos) => cos >= request.min_semantic_score,
824 None => true,
825 }
826 });
827 }
828
829 candidates.sort_by(|a, b| {
843 b.capsule
844 .score
845 .partial_cmp(&a.capsule.score)
846 .unwrap_or(Ordering::Equal)
847 .then_with(|| {
848 b.capsule
849 .freshness
850 .partial_cmp(&a.capsule.freshness)
851 .unwrap_or(Ordering::Equal)
852 })
853 .then_with(|| a.capsule.expansion_handle.cmp(&b.capsule.expansion_handle))
858 });
859
860 let embedding_mmr_ran = query_embedding.is_some() && !candidates.is_empty();
863 let candidates = if embedding_mmr_ran {
864 apply_candidate_mmr_diversity(candidates, 0.7)
865 } else {
866 candidates
867 };
868
869 let created_at_by_handle: std::collections::HashMap<String, String> =
873 if crate::ordering::is_ordering_query(&request.query) {
874 candidates
875 .iter()
876 .filter_map(|c| {
877 c.created_at
878 .clone()
879 .map(|ts| (c.capsule.expansion_handle.clone(), ts))
880 })
881 .collect()
882 } else {
883 std::collections::HashMap::new()
884 };
885
886 let mut capsules = candidates
887 .into_iter()
888 .map(|candidate| candidate.capsule)
889 .collect::<Vec<_>>();
890
891 if !embedding_mmr_ran {
894 capsules.sort_by(|left, right| {
895 right
896 .score
897 .partial_cmp(&left.score)
898 .unwrap_or(Ordering::Equal)
899 .then_with(|| {
900 right
901 .freshness
902 .partial_cmp(&left.freshness)
903 .unwrap_or(Ordering::Equal)
904 })
905 .then_with(|| left.expansion_handle.cmp(&right.expansion_handle))
907 });
908 }
909
910 let top_score = capsules.first().map(|c| c.score).unwrap_or(0.0);
913 if request.min_score > 0.0 && top_score < request.min_score {
914 return Ok(ContextBundle {
915 stage: request.stage,
916 budget_tokens: request.budget_tokens,
917 used_tokens: 0,
918 capsules: Vec::new(),
919 excluded: capsules,
920 skipped: true,
921 top_score,
922 evidence_coverage: 0.0,
924 uncovered_terms: Vec::new(),
925 chronological: false,
927 });
928 }
929
930 let capsules = apply_mmr_diversity(capsules, 0.7);
938
939 let capsule_budget = request.budget_tokens / 2;
940 let mut used_tokens = 0u32;
941 let mut included = Vec::new();
942 let mut excluded = Vec::new();
943
944 for capsule in capsules {
945 if request.max_capsules > 0 && included.len() >= request.max_capsules {
947 excluded.push(capsule);
948 continue;
949 }
950 if used_tokens.saturating_add(capsule.token_estimate) <= capsule_budget {
951 used_tokens += capsule.token_estimate;
952 included.push(capsule);
953 } else {
954 excluded.push(capsule);
955 }
956 }
957
958 let (coverage, uncovered_terms) = evidence_coverage(conn, &request.query, &included);
959
960 let chronological = !created_at_by_handle.is_empty();
967 let included = if chronological {
968 let dated = crate::ordering::render_chronologically(included, &created_at_by_handle);
969 used_tokens = dated.iter().map(|c| c.token_estimate).sum();
970 dated
971 } else {
972 included
973 };
974
975 Ok(ContextBundle {
976 stage: request.stage,
977 budget_tokens: request.budget_tokens,
978 used_tokens,
979 capsules: included,
980 excluded,
981 skipped: false,
982 top_score,
983 evidence_coverage: coverage,
984 uncovered_terms,
985 chronological,
986 })
987}
988
989pub fn search_memories_including_expired(
1001 conn: &Connection,
1002 limit: u32,
1003) -> KimetsuResult<Vec<ContextCapsule>> {
1004 let mut stmt = conn.prepare_cached(
1005 "
1006 SELECT memory_id, scope, kind, text, confidence, created_at,
1007 use_count, usefulness_score, valid_from, valid_to
1008 FROM memories
1009 WHERE invalidated_at IS NULL
1010 AND superseded_by IS NULL
1011 ORDER BY created_at DESC
1012 LIMIT ?1
1013 ",
1014 )?;
1015 let rows = stmt.query_map(params![limit], |row| {
1016 Ok((
1017 row.get::<_, String>(0)?,
1018 row.get::<_, String>(1)?,
1019 row.get::<_, String>(2)?,
1020 row.get::<_, String>(3)?,
1021 row.get::<_, f32>(4)?,
1022 row.get::<_, String>(5)?,
1023 row.get::<_, i64>(6)?,
1024 row.get::<_, f64>(7)?,
1025 row.get::<_, Option<String>>(8)?,
1026 row.get::<_, Option<String>>(9)?,
1027 ))
1028 })?;
1029 let now_utc = OffsetDateTime::now_utc();
1030 let now_rfc3339 = now_utc
1031 .format(&time::format_description::well_known::Rfc3339)
1032 .unwrap_or_default();
1033 let mut capsules = Vec::new();
1034 for row in rows {
1035 let (
1036 memory_id,
1037 scope,
1038 kind,
1039 text,
1040 confidence,
1041 created_at,
1042 _use_count,
1043 _usefulness,
1044 _valid_from,
1045 valid_to,
1046 ) = row?;
1047 let freshness = freshness(&created_at);
1048 let scope_weight = scope_weight(&scope);
1049 let suffix = if let Some(ref vt) = valid_to {
1051 if vt.as_str() < now_rfc3339.as_str() {
1052 format!(" [expired valid_to={vt}]")
1053 } else {
1054 format!(" [valid_to={vt}]")
1055 }
1056 } else {
1057 String::new()
1058 };
1059 capsules.push(ContextCapsule {
1060 id: new_id().to_string(),
1061 kind: "memory".to_string(),
1062 summary: format!("{scope}:{kind} - {text}{suffix}"),
1063 token_estimate: estimate_tokens(&text) + 8,
1064 expansion_handle: format!("memory:{memory_id}"),
1065 provenance: vec![ProvenanceRef {
1066 source: "Memory".to_string(),
1067 id: memory_id,
1068 excerpt: Some(excerpt(&text)),
1069 }],
1070 confidence,
1071 freshness,
1072 relevance: 0.0,
1073 scope_weight,
1074 score: 0.0,
1075 });
1076 }
1077 Ok(capsules)
1078}
1079
1080pub fn search_repo_files(
1081 conn: &Connection,
1082 repo_root: &str,
1083 query: &str,
1084 limit: u32,
1085) -> KimetsuResult<Vec<ContextCapsule>> {
1086 let candidates = repo_file_candidates(conn, repo_root, query, limit)?;
1087 let mut capsules = candidates
1088 .into_iter()
1089 .map(|mut candidate| {
1090 candidate.capsule.relevance = candidate.raw_relevance;
1091 candidate.capsule.score = candidate.raw_relevance;
1092 candidate.capsule
1093 })
1094 .collect::<Vec<_>>();
1095 capsules.sort_by(|left, right| {
1096 right
1097 .score
1098 .partial_cmp(&left.score)
1099 .unwrap_or(Ordering::Equal)
1100 .then_with(|| left.expansion_handle.cmp(&right.expansion_handle))
1101 });
1102 Ok(capsules)
1103}
1104
1105#[cfg(feature = "embeddings")]
1117fn memory_ann_candidates(
1118 conn: &Connection,
1119 qe: &QueryEmbedding,
1120 k: u32,
1121 query_tokens: &[String],
1122 half_life_days: f32,
1123) -> KimetsuResult<Vec<Candidate>> {
1124 let handle = crate::ann::handle_for_query(conn, qe.vector.len(), &qe.model_id)?;
1126 let hits = handle
1127 .read()
1128 .unwrap_or_else(|p| p.into_inner())
1129 .search(&qe.vector, k as usize)?;
1130 let knn_rowids: Vec<i64> = hits.into_iter().map(|(rowid, _dist)| rowid).collect();
1134 if knn_rowids.is_empty() {
1135 return Ok(Vec::new());
1136 }
1137
1138 let placeholders: String = knn_rowids
1140 .iter()
1141 .enumerate()
1142 .map(|(i, _)| format!("?{}", i + 1))
1143 .collect::<Vec<_>>()
1144 .join(", ");
1145 let sql = format!(
1146 "SELECT memory_id, scope, kind, text, confidence, created_at,
1147 use_count, usefulness_score, embedding, embedding_model,
1148 last_useful_at, provenance_snapshot_json
1149 FROM memories
1150 WHERE invalidated_at IS NULL
1151 AND superseded_by IS NULL
1152 AND (valid_to IS NULL OR valid_to > datetime('now'))
1153 AND embedding_model = ?{model_param}
1154 AND rowid IN ({placeholders})",
1155 model_param = knn_rowids.len() + 1
1156 );
1157 let mut stmt = conn.prepare(&sql)?;
1158 let mut params_vec: Vec<&dyn rusqlite::ToSql> = knn_rowids
1159 .iter()
1160 .map(|n| n as &dyn rusqlite::ToSql)
1161 .collect();
1162 params_vec.push(&qe.model_id);
1163 let rows_iter = stmt.query_map(params_vec.as_slice(), |row| {
1164 Ok((
1165 row.get::<_, String>(0)?,
1166 row.get::<_, String>(1)?,
1167 row.get::<_, String>(2)?,
1168 row.get::<_, String>(3)?,
1169 row.get::<_, f32>(4)?,
1170 row.get::<_, String>(5)?,
1171 row.get::<_, i64>(6)?,
1172 row.get::<_, f64>(7)?,
1173 row.get::<_, Option<Vec<u8>>>(8)?,
1174 row.get::<_, Option<String>>(9)?,
1175 row.get::<_, Option<String>>(10)?,
1176 row.get::<_, Option<String>>(11)?,
1177 ))
1178 })?;
1179
1180 let mut candidates = Vec::new();
1181 for row in rows_iter {
1182 let (
1183 memory_id,
1184 scope,
1185 kind,
1186 text,
1187 confidence,
1188 created_at,
1189 use_count,
1190 usefulness_score,
1191 embedding,
1192 embedding_model,
1193 last_useful_at,
1194 provenance_snapshot,
1195 ) = row?;
1196 let (cosine, row_vec) =
1197 compute_cosine_and_vec(Some(qe), embedding.as_deref(), embedding_model.as_deref());
1198 if let Some(candidate) = memory_row_to_candidate(
1199 query_tokens,
1200 memory_id,
1201 scope,
1202 kind,
1203 text,
1204 confidence,
1205 created_at,
1206 use_count,
1207 usefulness_score,
1208 last_useful_at,
1209 provenance_snapshot,
1210 half_life_days,
1211 None, cosine,
1213 row_vec,
1214 ) {
1215 candidates.push(candidate);
1216 }
1217 }
1218 Ok(candidates)
1219}
1220
1221pub(crate) fn memory_candidates_flat(
1227 conn: &Connection,
1228 query: &str,
1229 query_embedding: Option<&QueryEmbedding>,
1230 half_life_days: f32,
1231 fusion: crate::fusion::Fusion,
1232) -> KimetsuResult<Vec<Candidate>> {
1233 memory_candidates(conn, query, query_embedding, half_life_days, fusion)
1234}
1235
1236fn memory_candidates(
1243 conn: &Connection,
1244 query: &str,
1245 query_embedding: Option<&QueryEmbedding>,
1246 half_life_days: f32,
1247 #[cfg_attr(not(feature = "embeddings"), allow(unused_variables))] fusion: crate::fusion::Fusion,
1250) -> KimetsuResult<Vec<Candidate>> {
1251 let query_tokens = query_tokens(query);
1252
1253 #[cfg(feature = "embeddings")]
1258 if let Some(qe) = query_embedding {
1259 let fts_candidates = if let Some(fts_query) = fts_query(query) {
1261 memory_fts_candidates(
1262 conn,
1263 &query_tokens,
1264 &fts_query,
1265 80,
1266 Some(qe),
1267 half_life_days,
1268 )?
1269 } else {
1270 Vec::new()
1271 };
1272
1273 let ann_candidates = memory_ann_candidates(conn, qe, 80, &query_tokens, half_life_days)?;
1275
1276 return Ok(crate::fusion::fuse(
1278 fusion,
1279 vec![fts_candidates, ann_candidates],
1280 ));
1281 }
1282
1283 if let Some(fts_query) = fts_query(query) {
1285 let candidates = memory_fts_candidates(
1286 conn,
1287 &query_tokens,
1288 &fts_query,
1289 80,
1290 query_embedding,
1291 half_life_days,
1292 )?;
1293 if !candidates.is_empty() {
1294 return Ok(candidates);
1295 }
1296 }
1297
1298 latest_memory_candidates(conn, &query_tokens, 200, query_embedding, half_life_days)
1299}
1300
1301fn latest_memory_candidates(
1302 conn: &Connection,
1303 query_tokens: &[String],
1304 limit: u32,
1305 query_embedding: Option<&QueryEmbedding>,
1306 half_life_days: f32,
1307) -> KimetsuResult<Vec<Candidate>> {
1308 let mut stmt = conn.prepare_cached(
1319 "
1320 SELECT memory_id, scope, kind, text, confidence, created_at,
1321 use_count, usefulness_score, embedding, embedding_model,
1322 last_useful_at, provenance_snapshot_json
1323 FROM memories
1324 WHERE invalidated_at IS NULL
1325 AND superseded_by IS NULL
1326 AND (valid_to IS NULL OR valid_to > datetime('now'))
1327 ORDER BY created_at DESC
1328 LIMIT ?1
1329 ",
1330 )?;
1331
1332 let rows = stmt.query_map(params![limit], |row| {
1333 Ok((
1334 row.get::<_, String>(0)?,
1335 row.get::<_, String>(1)?,
1336 row.get::<_, String>(2)?,
1337 row.get::<_, String>(3)?,
1338 row.get::<_, f32>(4)?,
1339 row.get::<_, String>(5)?,
1340 row.get::<_, i64>(6)?,
1341 row.get::<_, f64>(7)?,
1342 row.get::<_, Option<Vec<u8>>>(8)?,
1343 row.get::<_, Option<String>>(9)?,
1344 row.get::<_, Option<String>>(10)?,
1345 row.get::<_, Option<String>>(11)?,
1346 ))
1347 })?;
1348
1349 let mut candidates = Vec::new();
1350 for row in rows {
1351 let (
1352 memory_id,
1353 scope,
1354 kind,
1355 text,
1356 confidence,
1357 created_at,
1358 use_count,
1359 usefulness_score,
1360 embedding,
1361 embedding_model,
1362 last_useful_at,
1363 provenance_snapshot,
1364 ) = row?;
1365 let (cosine, row_vec) = compute_cosine_and_vec(
1366 query_embedding,
1367 embedding.as_deref(),
1368 embedding_model.as_deref(),
1369 );
1370 if let Some(candidate) = memory_row_to_candidate(
1371 query_tokens,
1372 memory_id,
1373 scope,
1374 kind,
1375 text,
1376 confidence,
1377 created_at,
1378 use_count,
1379 usefulness_score,
1380 last_useful_at,
1381 provenance_snapshot,
1382 half_life_days,
1383 None,
1384 cosine,
1385 row_vec,
1386 ) {
1387 candidates.push(candidate);
1388 }
1389 }
1390 Ok(candidates)
1391}
1392
1393fn memory_fts_candidates(
1394 conn: &Connection,
1395 query_tokens: &[String],
1396 fts_query: &str,
1397 limit: u32,
1398 query_embedding: Option<&QueryEmbedding>,
1399 half_life_days: f32,
1400) -> KimetsuResult<Vec<Candidate>> {
1401 let mut stmt = conn.prepare_cached(
1402 "
1403 SELECT m.memory_id, m.scope, m.kind, m.text, m.confidence, m.created_at,
1404 m.use_count, m.usefulness_score, bm25(memories_fts) AS rank,
1405 m.embedding, m.embedding_model, m.last_useful_at,
1406 m.provenance_snapshot_json
1407 FROM memories_fts
1408 JOIN memories m
1409 ON m.memory_id = memories_fts.memory_id
1410 WHERE m.invalidated_at IS NULL
1411 AND m.superseded_by IS NULL
1412 AND (m.valid_to IS NULL OR m.valid_to > datetime('now'))
1413 AND memories_fts MATCH ?1
1414 ORDER BY rank
1415 LIMIT ?2
1416 ",
1417 )?;
1418
1419 let rows = stmt.query_map(params![fts_query, limit], |row| {
1420 Ok((
1421 row.get::<_, String>(0)?,
1422 row.get::<_, String>(1)?,
1423 row.get::<_, String>(2)?,
1424 row.get::<_, String>(3)?,
1425 row.get::<_, f32>(4)?,
1426 row.get::<_, String>(5)?,
1427 row.get::<_, i64>(6)?,
1428 row.get::<_, f64>(7)?,
1429 row.get::<_, f64>(8)?,
1430 row.get::<_, Option<Vec<u8>>>(9)?,
1431 row.get::<_, Option<String>>(10)?,
1432 row.get::<_, Option<String>>(11)?,
1433 row.get::<_, Option<String>>(12)?,
1434 ))
1435 })?;
1436
1437 let mut candidates = Vec::new();
1438 for row in rows {
1439 let (
1440 memory_id,
1441 scope,
1442 kind,
1443 text,
1444 confidence,
1445 created_at,
1446 use_count,
1447 usefulness_score,
1448 rank,
1449 embedding,
1450 embedding_model,
1451 last_useful_at,
1452 provenance_snapshot,
1453 ) = row?;
1454 let fts_relevance = (-rank as f32).max(0.0);
1455 let (cosine, row_vec) = compute_cosine_and_vec(
1456 query_embedding,
1457 embedding.as_deref(),
1458 embedding_model.as_deref(),
1459 );
1460 if let Some(candidate) = memory_row_to_candidate(
1461 query_tokens,
1462 memory_id,
1463 scope,
1464 kind,
1465 text,
1466 confidence,
1467 created_at,
1468 use_count,
1469 usefulness_score,
1470 last_useful_at,
1471 provenance_snapshot,
1472 half_life_days,
1473 Some(fts_relevance),
1474 cosine,
1475 row_vec,
1476 ) {
1477 candidates.push(candidate);
1478 }
1479 }
1480 Ok(candidates)
1481}
1482
1483fn compute_cosine_and_vec(
1507 query_embedding: Option<&QueryEmbedding>,
1508 row_bytes: Option<&[u8]>,
1509 row_model: Option<&str>,
1510) -> (Option<f32>, Option<Vec<f32>>) {
1511 let q = match query_embedding {
1512 Some(q) => q,
1513 None => return (None, None),
1514 };
1515 let bytes = match row_bytes {
1516 Some(b) => b,
1517 None => return (None, None),
1518 };
1519 let model = match row_model {
1520 Some(m) => m,
1521 None => return (None, None),
1522 };
1523 if model != q.model_id {
1524 return (None, None);
1525 }
1526 let row_vec = match decode_embedding(bytes, Some(q.vector.len())) {
1527 Ok(v) => v,
1528 Err(_) => return (None, None),
1529 };
1530 let score = cosine_similarity(&q.vector, &row_vec);
1531 (Some(score), Some(row_vec))
1532}
1533
1534#[allow(clippy::too_many_arguments)]
1535fn memory_row_to_candidate(
1536 query_tokens: &[String],
1537 memory_id: String,
1538 scope: String,
1539 kind: String,
1540 text: String,
1541 confidence: f32,
1542 created_at: String,
1543 use_count: i64,
1544 usefulness_score: f64,
1545 last_useful_at: Option<String>,
1546 provenance_snapshot: Option<String>,
1549 half_life_days: f32,
1550 raw_relevance_override: Option<f32>,
1551 cosine_score: Option<f32>,
1552 row_embedding: Option<Vec<f32>>,
1557) -> Option<Candidate> {
1558 let lexical = lexical_relevance(query_tokens, &format!("{kind} {text}"));
1559 let lexical_term = raw_relevance_override.unwrap_or(lexical).max(lexical);
1560
1561 let raw_relevance = match cosine_score {
1573 Some(c) => {
1574 let normalized_cos = ((c + 1.0) * 0.5).clamp(0.0, 1.0);
1575 (1.0 - DEFAULT_HYBRID_ALPHA) * lexical_term + DEFAULT_HYBRID_ALPHA * normalized_cos
1576 }
1577 None => lexical_term,
1578 };
1579
1580 if raw_relevance <= 0.0 && !query_tokens.is_empty() {
1586 return None;
1587 }
1588
1589 let freshness = freshness(&created_at);
1590 let scope_weight = scope_weight(&scope);
1591 let raw_multiplier = usefulness_multiplier(usefulness_score as f32, use_count as u32);
1597 let decay = usefulness_decay(last_useful_at.as_deref(), &created_at, half_life_days);
1598 let multiplier = 1.0 + (raw_multiplier - 1.0) * decay;
1599 let biased_relevance = apply_usefulness_boost(raw_relevance, multiplier);
1600
1601 let provenance =
1610 crate::trust::Provenance::from_snapshot(provenance_snapshot.as_deref().unwrap_or("{}"));
1611 let trusted_relevance =
1612 biased_relevance * crate::trust::trust_multiplier(provenance, last_useful_at.is_some());
1613
1614 Some(Candidate {
1615 raw_relevance: trusted_relevance,
1616 embedding: row_embedding,
1617 cosine: cosine_score,
1618 created_at: Some(created_at),
1619 capsule: ContextCapsule {
1620 id: new_id().to_string(),
1621 kind: "memory".to_string(),
1622 summary: format!("{scope}:{kind} - {text}"),
1623 token_estimate: estimate_tokens(&text) + 8,
1624 expansion_handle: format!("memory:{memory_id}"),
1625 provenance: vec![ProvenanceRef {
1626 source: "Memory".to_string(),
1627 id: memory_id,
1628 excerpt: Some(excerpt(&text)),
1629 }],
1630 confidence,
1631 freshness,
1632 relevance: 0.0,
1633 scope_weight,
1634 score: 0.0,
1635 },
1636 })
1637}
1638
1639pub(crate) fn usefulness_decay(
1663 last_useful_at: Option<&str>,
1664 created_at: &str,
1665 half_life_days: f32,
1666) -> f32 {
1667 if half_life_days <= 0.0 {
1668 return 1.0;
1669 }
1670 let reference = last_useful_at.unwrap_or(created_at);
1671 let Ok(reference_ts) =
1672 OffsetDateTime::parse(reference, &time::format_description::well_known::Rfc3339)
1673 else {
1674 return 1.0;
1675 };
1676 let age = OffsetDateTime::now_utc() - reference_ts;
1677 let age_days = (age.whole_seconds().max(0) as f32) / 86_400.0;
1678 let exponent = -std::f32::consts::LN_2 * age_days / half_life_days;
1679 exponent.exp().clamp(0.0, 1.0)
1680}
1681
1682pub(crate) use crate::scoring::USEFULNESS_BOOST_CAP;
1685
1686pub(crate) fn apply_usefulness_boost(raw_relevance: f32, multiplier: f32) -> f32 {
1689 if multiplier <= 1.0 {
1690 return raw_relevance * multiplier;
1691 }
1692 (raw_relevance * multiplier).min(raw_relevance + USEFULNESS_BOOST_CAP)
1693}
1694
1695pub(crate) fn usefulness_multiplier(usefulness_score: f32, use_count: u32) -> f32 {
1700 use crate::scoring::{FULL_CONFIDENCE_USES, MULTIPLIER_MAX, MULTIPLIER_MIN};
1708 if use_count == 0 {
1709 return 1.0;
1710 }
1711 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);
1714 let confidence = (use_count as f32 / FULL_CONFIDENCE_USES as f32).min(1.0);
1715 1.0 * (1.0 - confidence) + full_multiplier * confidence
1716}
1717
1718fn repo_file_candidates(
1719 conn: &Connection,
1720 repo_root: &str,
1721 query: &str,
1722 limit: u32,
1723) -> KimetsuResult<Vec<Candidate>> {
1724 let Some(fts_query) = fts_query(query) else {
1725 return Ok(Vec::new());
1726 };
1727
1728 let mut stmt = conn.prepare_cached(
1729 "
1730 SELECT path, snippet, language_guess, bm25(repo_files_fts) AS rank
1731 FROM repo_files_fts
1732 WHERE repo_root = ?1 AND repo_files_fts MATCH ?2
1733 ORDER BY rank
1734 LIMIT ?3
1735 ",
1736 )?;
1737
1738 let rows = stmt.query_map(params![repo_root, fts_query, limit], |row| {
1739 Ok((
1740 row.get::<_, String>(0)?,
1741 row.get::<_, String>(1)?,
1742 row.get::<_, String>(2)?,
1743 row.get::<_, f64>(3)?,
1744 ))
1745 })?;
1746
1747 let mut candidates = Vec::new();
1748 for row in rows {
1749 let (path, snippet, language, rank) = row?;
1750 let raw_relevance = (-rank as f32).max(0.0);
1751 let summary = format!("{path} ({language}) - {}", excerpt(&snippet));
1752 let token_estimate = estimate_tokens(&summary) + 8;
1753 candidates.push(Candidate {
1754 raw_relevance,
1755 embedding: None,
1756 cosine: None,
1757 created_at: None,
1759 capsule: ContextCapsule {
1760 id: new_id().to_string(),
1761 kind: "repo_file".to_string(),
1762 summary,
1763 token_estimate,
1764 expansion_handle: format!("file:{path}"),
1765 provenance: vec![ProvenanceRef {
1766 source: "RepoFile".to_string(),
1767 id: path.clone(),
1768 excerpt: Some(excerpt(&snippet)),
1769 }],
1770 confidence: 0.9,
1771 freshness: 1.0,
1772 relevance: 0.0,
1773 scope_weight: 0.9,
1774 score: 0.0,
1775 },
1776 });
1777 }
1778 Ok(candidates)
1779}
1780
1781fn manifest_candidates(
1782 conn: &Connection,
1783 repo_root: &str,
1784 query: &str,
1785) -> KimetsuResult<Vec<Candidate>> {
1786 if let Some(fts_query) = fts_query(query) {
1787 let candidates = manifest_fts_candidates(conn, repo_root, &fts_query, 30)?;
1788 if !candidates.is_empty() {
1789 return Ok(candidates);
1790 }
1791 }
1792
1793 let query_tokens = query_tokens(query);
1794 let mut stmt = conn.prepare_cached(
1795 "
1796 SELECT manifest_path, manifest_kind, parsed_summary_json
1797 FROM repo_manifests
1798 WHERE repo_root = ?1
1799 ORDER BY manifest_path
1800 ",
1801 )?;
1802
1803 let rows = stmt.query_map(params![repo_root], |row| {
1804 Ok((
1805 row.get::<_, String>(0)?,
1806 row.get::<_, String>(1)?,
1807 row.get::<_, String>(2)?,
1808 ))
1809 })?;
1810
1811 let mut candidates = Vec::new();
1812 for row in rows {
1813 let (path, kind, summary_json) = row?;
1814 let raw_relevance =
1815 lexical_relevance(&query_tokens, &format!("{path} {kind} {summary_json}"));
1816 if raw_relevance <= 0.0 && !query_tokens.is_empty() {
1817 continue;
1818 }
1819 let summary = format!("{path} manifest ({kind})");
1820 let token_estimate = estimate_tokens(&summary) + 8;
1821 candidates.push(Candidate {
1822 raw_relevance,
1823 embedding: None,
1824 cosine: None,
1825 created_at: None,
1827 capsule: ContextCapsule {
1828 id: new_id().to_string(),
1829 kind: "repo_manifest".to_string(),
1830 summary,
1831 token_estimate,
1832 expansion_handle: format!("file:{path}"),
1833 provenance: vec![ProvenanceRef {
1834 source: "Manifest".to_string(),
1835 id: path,
1836 excerpt: Some(excerpt(&summary_json)),
1837 }],
1838 confidence: 0.95,
1839 freshness: 1.0,
1840 relevance: 0.0,
1841 scope_weight: 0.9,
1842 score: 0.0,
1843 },
1844 });
1845 }
1846 Ok(candidates)
1847}
1848
1849fn manifest_fts_candidates(
1850 conn: &Connection,
1851 repo_root: &str,
1852 fts_query: &str,
1853 limit: u32,
1854) -> KimetsuResult<Vec<Candidate>> {
1855 let mut stmt = conn.prepare_cached(
1856 "
1857 SELECT manifest_path, manifest_kind, parsed_summary_json,
1858 bm25(repo_manifests_fts) AS rank
1859 FROM repo_manifests_fts
1860 WHERE repo_root = ?1 AND repo_manifests_fts MATCH ?2
1861 ORDER BY rank
1862 LIMIT ?3
1863 ",
1864 )?;
1865
1866 let rows = stmt.query_map(params![repo_root, fts_query, limit], |row| {
1867 Ok((
1868 row.get::<_, String>(0)?,
1869 row.get::<_, String>(1)?,
1870 row.get::<_, String>(2)?,
1871 row.get::<_, f64>(3)?,
1872 ))
1873 })?;
1874
1875 let mut candidates = Vec::new();
1876 for row in rows {
1877 let (path, kind, summary_json, rank) = row?;
1878 let raw_relevance = (-rank as f32).max(0.0);
1879 let summary = format!("{path} manifest ({kind})");
1880 let token_estimate = estimate_tokens(&summary) + 8;
1881 candidates.push(Candidate {
1882 raw_relevance,
1883 embedding: None,
1884 cosine: None,
1885 created_at: None,
1887 capsule: ContextCapsule {
1888 id: new_id().to_string(),
1889 kind: "repo_manifest".to_string(),
1890 summary,
1891 token_estimate,
1892 expansion_handle: format!("file:{path}"),
1893 provenance: vec![ProvenanceRef {
1894 source: "Manifest".to_string(),
1895 id: path,
1896 excerpt: Some(excerpt(&summary_json)),
1897 }],
1898 confidence: 0.95,
1899 freshness: 1.0,
1900 relevance: 0.0,
1901 scope_weight: 0.9,
1902 score: 0.0,
1903 },
1904 });
1905 }
1906 Ok(candidates)
1907}
1908
1909#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1941pub enum Normalization {
1942 #[default]
1944 PerKind,
1945 Global,
1947}
1948
1949impl Normalization {
1950 pub fn from_config(value: &str) -> Self {
1954 match value.trim().to_ascii_lowercase().as_str() {
1955 "global" => Self::Global,
1956 _ => Self::PerKind,
1957 }
1958 }
1959}
1960
1961fn normalize_and_score(
1962 candidates: &mut [Candidate],
1963 weights: StageWeights,
1964 normalization: Normalization,
1965) {
1966 let mut max_by_kind = HashMap::<String, f32>::new();
1969 let bucket = |candidate: &Candidate| match normalization {
1970 Normalization::PerKind => candidate.capsule.kind.clone(),
1971 Normalization::Global => String::new(),
1972 };
1973 for candidate in candidates.iter() {
1974 max_by_kind
1975 .entry(bucket(candidate))
1976 .and_modify(|max| *max = (*max).max(candidate.raw_relevance))
1977 .or_insert(candidate.raw_relevance);
1978 }
1979
1980 for candidate in candidates {
1981 let max = max_by_kind.get(&bucket(candidate)).copied().unwrap_or(0.0);
1982 let relevance = if max <= f32::EPSILON {
1983 if candidate.raw_relevance > 0.0 {
1984 1.0
1985 } else {
1986 0.0
1987 }
1988 } else {
1989 (candidate.raw_relevance / max).clamp(0.0, 1.0)
1990 };
1991 candidate.capsule.relevance = relevance;
1992 candidate.capsule.score = weights.relevance * relevance
1993 + weights.confidence * candidate.capsule.confidence
1994 + weights.freshness * candidate.capsule.freshness
1995 + weights.scope * candidate.capsule.scope_weight;
1996 }
1997}
1998
1999fn weights_for_stage(weights: &BrokerWeights, stage: &str) -> StageWeights {
2000 match stage {
2001 "localization" => weights.localization.clone(),
2002 "patch_plan" => weights.patch_plan.clone(),
2003 "verification" => weights.verification.clone(),
2004 "review" => weights.review.clone(),
2005 _ => None,
2006 }
2007 .unwrap_or(StageWeights {
2008 relevance: weights.relevance,
2009 confidence: weights.confidence,
2010 freshness: weights.freshness,
2011 scope: weights.scope,
2012 })
2013}
2014
2015pub(crate) fn scope_weight_pub(scope: &str) -> f32 {
2018 scope_weight(scope)
2019}
2020
2021fn scope_weight(scope: &str) -> f32 {
2022 match scope.parse::<MemoryScope>() {
2023 Ok(MemoryScope::Run) => 1.0,
2024 Ok(MemoryScope::Repo) => 0.9,
2025 Ok(MemoryScope::Project) => 0.7,
2026 Ok(MemoryScope::GlobalUser) => 0.5,
2027 Err(_) => 0.3,
2028 }
2029}
2030
2031pub(crate) fn freshness_pub(created_at: &str) -> f32 {
2034 freshness(created_at)
2035}
2036
2037fn freshness(created_at: &str) -> f32 {
2038 let Ok(created_at) =
2039 OffsetDateTime::parse(created_at, &time::format_description::well_known::Rfc3339)
2040 else {
2041 return 0.5;
2042 };
2043 let age = OffsetDateTime::now_utc() - created_at;
2044 let age_days = age.whole_seconds().max(0) as f32 / 86_400.0;
2045 (-age_days / 30.0).exp().clamp(0.0, 1.0)
2046}
2047
2048const SEMANTIC_KEEP_COSINE: f32 = 0.20;
2053
2054const STOPWORDS: &[&str] = &[
2059 "the", "and", "for", "are", "but", "not", "you", "your", "with", "this", "that", "these",
2060 "those", "from", "into", "about", "what", "whats", "which", "who", "whom", "how", "why",
2061 "when", "where", "can", "could", "would", "should", "will", "shall", "does", "did", "was",
2062 "were", "been", "being", "have", "has", "had", "its", "it", "is", "as", "at", "by", "of", "to",
2063 "in", "on", "or", "an", "be", "do", "me", "my", "we", "us", "our", "im", "ive", "let", "lets",
2064 "please", "tell", "give", "show", "want", "need", "get", "got", "use", "using", "there",
2065 "their", "they", "them", "then", "than", "some", "any", "all", "more", "most", "such", "via",
2066 "per",
2067 "during", "while", "until", "unless", "before", "after", "again", "against", "above", "below",
2076 "between", "through", "under", "over", "because", "also", "just", "only", "very", "much",
2077 "many", "each", "both", "same", "other", "another", "always", "never", "still", "even", "ever",
2078 "every", "first", "found", "thing", "things", "value", "default", "if", "so", "up", "out",
2079 "off", "down", "no", "yes",
2080];
2081
2082fn content_tokens(query: &str) -> Vec<String> {
2087 let mut seen = std::collections::HashSet::new();
2088 query
2089 .split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
2090 .map(str::trim)
2091 .filter(|part| part.len() >= 2)
2092 .map(str::to_ascii_lowercase)
2093 .filter(|t| !STOPWORDS.contains(&t.as_str()))
2094 .map(|t| light_stem(&t).to_string())
2097 .filter(|t| seen.insert(t.clone()))
2098 .collect()
2099}
2100
2101fn corpus_token_idf(conn: &Connection, tokens: &[String]) -> KimetsuResult<HashMap<String, f32>> {
2119 let mut idf = HashMap::new();
2120 let n: i64 = conn
2121 .query_row(
2122 "SELECT COUNT(*) FROM memories WHERE invalidated_at IS NULL",
2123 [],
2124 |row| row.get(0),
2125 )
2126 .unwrap_or(0);
2127 if n == 0 {
2128 return Ok(idf);
2129 }
2130 let mut stmt = conn.prepare_cached(
2131 "SELECT COUNT(*) FROM memories \
2132 WHERE invalidated_at IS NULL AND lower(text) LIKE ?1 ESCAPE '\\'",
2133 )?;
2134 for token in tokens {
2135 let pattern = format!("%{}%", escape_like(token));
2136 let df: i64 = stmt
2137 .query_row(params![pattern], |row| row.get(0))
2138 .unwrap_or(0);
2139 let weight = if df == 0 {
2141 0.0
2142 } else {
2143 (((n + 1) as f32) / ((df + 1) as f32)).ln().max(0.0)
2144 };
2145 idf.insert(token.clone(), weight);
2146 }
2147 Ok(idf)
2148}
2149
2150fn escape_like(token: &str) -> String {
2153 token
2154 .replace('\\', "\\\\")
2155 .replace('%', "\\%")
2156 .replace('_', "\\_")
2157}
2158
2159fn weighted_coverage(content: &[String], idf: &HashMap<String, f32>, summary: &str) -> f32 {
2166 let haystack = summary.to_ascii_lowercase();
2167 let mut total = 0.0f32;
2168 let mut hit = 0.0f32;
2169 for token in content {
2170 let weight = idf.get(token).copied().unwrap_or(0.0);
2171 total += weight;
2172 if weight > 0.0 && haystack.contains(token.as_str()) {
2173 hit += weight;
2174 }
2175 }
2176 if total <= f32::EPSILON {
2177 0.0
2178 } else {
2179 (hit / total).clamp(0.0, 1.0)
2180 }
2181}
2182
2183fn light_stem(token: &str) -> &str {
2208 let mut stem = token;
2209 for suffix in ["ing", "ed", "es", "s"] {
2210 if let Some(stripped) = token.strip_suffix(suffix)
2211 && stripped.len() >= 4
2212 {
2213 stem = stripped;
2214 break;
2215 }
2216 }
2217 if stem.len() >= 5
2218 && let Some(trimmed) = stem.strip_suffix('y').or_else(|| stem.strip_suffix('i'))
2219 && trimmed
2220 .chars()
2221 .next_back()
2222 .is_some_and(|c| !matches!(c, 'a' | 'e' | 'i' | 'o' | 'u'))
2223 {
2224 return trimmed;
2225 }
2226 stem
2227}
2228
2229fn query_tokens(query: &str) -> Vec<String> {
2230 let mut tokens: Vec<String> = query
2231 .split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
2232 .map(str::trim)
2233 .filter(|part| part.len() >= 2)
2234 .map(str::to_ascii_lowercase)
2235 .map(|t| light_stem(&t).to_string())
2236 .collect();
2237 let lower = query.to_ascii_lowercase();
2244 for (triggers, expansions) in CLASS_HINTS.iter() {
2245 if triggers.iter().any(|t| lower.contains(t)) {
2246 tokens.extend(expansions.iter().map(|e| e.to_string()));
2247 }
2248 }
2249 tokens
2250}
2251
2252const CLASS_HINTS: &[(&[&str], &[&str])] = &[
2260 (
2261 &[
2262 "build",
2263 "compile",
2264 "make",
2265 "cargo",
2266 "cmake",
2267 "configure",
2268 "install",
2269 "train",
2270 "benchmark",
2271 "test suite",
2272 "ray trace",
2273 "render",
2274 ],
2275 &[
2276 "shell_background",
2277 "shell_status",
2278 "shell_output",
2279 "shell_stop",
2280 "long_running",
2281 ],
2282 ),
2283 (
2284 &[
2285 "edit", "modify", "change", "fix", "update", "patch", "refactor", "rename",
2286 ],
2287 &["edit_file", "apply_patch", "old_string", "new_string"],
2288 ),
2289 (
2290 &[
2291 "read", "inspect", "review", "analyze", "examine", "view", "show",
2292 ],
2293 &["read_file", "offset", "limit", "multi_read"],
2294 ),
2295 (
2296 &["find", "locate", "search", "look up", "discover", "list"],
2297 &["glob", "search_files", "list_files"],
2298 ),
2299 (
2300 &["plan", "step", "checklist", "todo", "task list", "phase"],
2301 &["plan", "todos"],
2302 ),
2303 (
2304 &[
2305 "verify",
2306 "check",
2307 "ensure",
2308 "validate",
2309 "pass test",
2310 "verifier",
2311 ],
2312 &["finish", "verifier", "verification"],
2313 ),
2314 (
2315 &[
2316 "image",
2317 "png",
2318 "jpeg",
2319 "jpg",
2320 "pdf",
2321 "diagram",
2322 "screenshot",
2323 ],
2324 &["view_image", "base64", "sha256"],
2325 ),
2326 (&["delete", "remove", "rm "], &["delete_file", "recursive"]),
2327 (&["rename", "move file", "mv "], &["move_file"]),
2328];
2329
2330fn capsule_matches_kind(capsule: &ContextCapsule, wanted: &str) -> bool {
2335 if capsule.kind == wanted {
2336 return true;
2337 }
2338 if capsule.kind == "memory"
2339 && let Some((prefix, _)) = capsule.summary.split_once(" - ")
2340 && let Some((_scope, mkind)) = prefix.split_once(':')
2341 {
2342 return mkind == wanted;
2343 }
2344 false
2345}
2346
2347pub(crate) fn fts_query(query: &str) -> Option<String> {
2348 let tokens = query_tokens(query);
2349 if tokens.is_empty() {
2350 return None;
2351 }
2352 Some(
2353 tokens
2354 .into_iter()
2355 .take(12)
2356 .map(|token| format!("{token}*"))
2357 .collect::<Vec<_>>()
2358 .join(" OR "),
2359 )
2360}
2361
2362fn apply_candidate_mmr_diversity(mut sorted: Vec<Candidate>, lambda: f32) -> Vec<Candidate> {
2382 if sorted.len() <= 1 {
2383 return sorted;
2384 }
2385 let summaries: Vec<std::collections::HashSet<String>> = sorted
2387 .iter()
2388 .map(|c| summary_token_set(&c.capsule.summary))
2389 .collect();
2390
2391 let mut picked_indices: Vec<usize> = Vec::with_capacity(sorted.len());
2392 let mut remaining: Vec<usize> = (0..sorted.len()).collect();
2393
2394 picked_indices.push(remaining.remove(0));
2396
2397 while !remaining.is_empty() {
2398 let mut best_idx_in_remaining = 0;
2399 let mut best_score = f32::MIN;
2400
2401 for (i, &cand) in remaining.iter().enumerate() {
2402 let mut max_overlap = 0.0f32;
2403 for &p in &picked_indices {
2404 let same_kind = sorted[cand].capsule.kind == sorted[p].capsule.kind;
2407 let raw_overlap = candidate_pair_overlap(
2408 &sorted[cand],
2409 &sorted[p],
2410 &summaries[cand],
2411 &summaries[p],
2412 );
2413 let overlap = if same_kind {
2414 raw_overlap
2415 } else {
2416 raw_overlap * 0.5
2417 };
2418 if overlap > max_overlap {
2419 max_overlap = overlap;
2420 }
2421 }
2422 let mmr = lambda * sorted[cand].capsule.score - (1.0 - lambda) * max_overlap;
2423 if mmr > best_score {
2424 best_score = mmr;
2425 best_idx_in_remaining = i;
2426 }
2427 }
2428 picked_indices.push(remaining.remove(best_idx_in_remaining));
2429 }
2430
2431 let mut taken: Vec<Option<Candidate>> = sorted.drain(..).map(Some).collect();
2433 let mut out = Vec::with_capacity(taken.len());
2434 for idx in picked_indices {
2435 if let Some(c) = taken[idx].take() {
2436 out.push(c);
2437 }
2438 }
2439 out
2440}
2441
2442fn candidate_pair_overlap(
2449 a: &Candidate,
2450 b: &Candidate,
2451 tokens_a: &std::collections::HashSet<String>,
2452 tokens_b: &std::collections::HashSet<String>,
2453) -> f32 {
2454 if let (Some(va), Some(vb)) = (a.embedding.as_deref(), b.embedding.as_deref()) {
2455 cosine_similarity(va, vb).max(0.0)
2460 } else {
2461 jaccard(tokens_a, tokens_b)
2462 }
2463}
2464
2465fn apply_mmr_diversity(mut sorted: Vec<ContextCapsule>, lambda: f32) -> Vec<ContextCapsule> {
2477 if sorted.len() <= 1 {
2478 return sorted;
2479 }
2480 let summaries: Vec<std::collections::HashSet<String>> = sorted
2482 .iter()
2483 .map(|c| summary_token_set(&c.summary))
2484 .collect();
2485 let mut picked_indices: Vec<usize> = Vec::with_capacity(sorted.len());
2486 let mut remaining: Vec<usize> = (0..sorted.len()).collect();
2487
2488 picked_indices.push(remaining.remove(0));
2490
2491 while !remaining.is_empty() {
2492 let mut best_idx_in_remaining = 0;
2493 let mut best_score = f32::MIN;
2494 for (i, &cand) in remaining.iter().enumerate() {
2495 let mut max_overlap = 0.0f32;
2496 for &p in &picked_indices {
2497 let raw = jaccard(&summaries[cand], &summaries[p]);
2498 let overlap = if sorted[cand].kind == sorted[p].kind {
2499 raw
2500 } else {
2501 raw * 0.5
2504 };
2505 if overlap > max_overlap {
2506 max_overlap = overlap;
2507 }
2508 }
2509 let mmr = lambda * sorted[cand].score - (1.0 - lambda) * max_overlap;
2510 if mmr > best_score {
2511 best_score = mmr;
2512 best_idx_in_remaining = i;
2513 }
2514 }
2515 picked_indices.push(remaining.remove(best_idx_in_remaining));
2516 }
2517 let mut out = Vec::with_capacity(sorted.len());
2519 let mut taken: Vec<Option<ContextCapsule>> = sorted.drain(..).map(Some).collect();
2521 for idx in picked_indices {
2522 if let Some(c) = taken[idx].take() {
2523 out.push(c);
2524 }
2525 }
2526 out
2527}
2528
2529fn summary_token_set(s: &str) -> std::collections::HashSet<String> {
2530 s.split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
2531 .filter(|t| t.len() >= 3)
2532 .map(str::to_ascii_lowercase)
2533 .collect()
2534}
2535
2536fn jaccard(a: &std::collections::HashSet<String>, b: &std::collections::HashSet<String>) -> f32 {
2537 if a.is_empty() && b.is_empty() {
2538 return 0.0;
2539 }
2540 let intersection = a.intersection(b).count();
2541 let union = a.union(b).count();
2542 intersection as f32 / union.max(1) as f32
2543}
2544
2545fn lexical_relevance(tokens: &[String], haystack: &str) -> f32 {
2546 if tokens.is_empty() {
2547 return 0.0;
2548 }
2549 let haystack = haystack.to_ascii_lowercase();
2550 let matches = tokens
2551 .iter()
2552 .filter(|token| haystack.contains(token.as_str()))
2553 .count();
2554 matches as f32 / tokens.len() as f32
2555}
2556
2557pub fn estimate_tokens(text: &str) -> u32 {
2558 ((text.split_whitespace().count() as f32) * 1.33).ceil() as u32
2559}
2560
2561pub fn compress_for_render(summary: &str, max_sentences: usize) -> String {
2584 if max_sentences == 0 {
2585 return summary.to_string();
2586 }
2587
2588 let text = if let Some(rest) = summary.strip_prefix('[') {
2590 if let Some(idx) = rest.find(']') {
2592 rest[idx + 1..].trim_start()
2593 } else {
2594 summary
2595 }
2596 } else {
2597 summary
2598 };
2599
2600 let text = if let Some(idx) = text.rfind('(') {
2602 let candidate = text[..idx].trim_end();
2603 let inner = &text[idx + 1..];
2606 if inner.contains(':') && inner.trim_end().ends_with(')') {
2607 candidate
2608 } else {
2609 text
2610 }
2611 } else {
2612 text
2613 };
2614
2615 let (scope_prefix, body) = if let Some(dash_pos) = text.find(" - ") {
2617 let prefix_candidate = &text[..dash_pos];
2618 if !prefix_candidate.contains(' ') && prefix_candidate.contains(':') {
2620 let body_start = dash_pos + 3; (&text[..body_start], &text[body_start..])
2622 } else {
2623 ("", text)
2624 }
2625 } else {
2626 ("", text)
2627 };
2628
2629 let compressed_body = cap_sentences(body, max_sentences);
2631
2632 let result = if scope_prefix.is_empty() {
2634 compressed_body.to_string()
2635 } else {
2636 format!("{scope_prefix}{compressed_body}")
2637 };
2638
2639 if result.trim().is_empty() {
2640 summary.to_string()
2641 } else {
2642 result
2643 }
2644}
2645
2646fn cap_sentences(text: &str, n: usize) -> &str {
2650 let bytes = text.as_bytes();
2651 let len = bytes.len();
2652 let mut count = 0;
2653 let mut i = 0;
2654 while i < len {
2655 if bytes[i] == b'.' {
2657 let next = i + 1;
2658 if next < len && (bytes[next] == b' ' || bytes[next] == b'\n') {
2659 count += 1;
2660 if count >= n {
2661 return text[..=i].trim_end();
2663 }
2664 }
2665 }
2666 i += 1;
2667 }
2668 text.trim_end()
2670}
2671
2672pub(crate) fn excerpt_pub(text: &str) -> String {
2675 excerpt(text)
2676}
2677
2678fn excerpt(text: &str) -> String {
2679 let value = one_line(text);
2680 value.chars().take(256).collect()
2681}
2682
2683fn one_line(text: &str) -> String {
2684 text.split_whitespace().collect::<Vec<_>>().join(" ")
2685}
2686
2687const FILE_EXPAND_CAP_BYTES: usize = 2048;
2694
2695pub fn resolve_capsule(
2707 conn: &Connection,
2708 repo_root: &std::path::Path,
2709 handle: &str,
2710) -> kimetsu_core::KimetsuResult<String> {
2711 if let Some(memory_id) = handle.strip_prefix("memory:") {
2712 let mut stmt = conn.prepare_cached(
2714 "SELECT text FROM memories WHERE memory_id = ? AND invalidated_at IS NULL",
2715 )?;
2716 let text: Option<String> = stmt
2717 .query_row(rusqlite::params![memory_id], |row| row.get(0))
2718 .optional()?;
2719 match text {
2720 Some(t) => Ok(t),
2721 None => {
2722 Err(format!("expand_capsule: no active memory found for handle `{handle}`").into())
2723 }
2724 }
2725 } else if let Some(rel_path) = handle.strip_prefix("file:") {
2726 let path = std::path::Path::new(rel_path);
2731 if path.is_absolute() {
2732 return Err(format!(
2733 "expand_capsule: `{handle}` is an absolute path — only repo-relative paths are supported"
2734 )
2735 .into());
2736 }
2737 for component in path.components() {
2738 match component {
2739 std::path::Component::ParentDir => {
2740 return Err(format!(
2741 "expand_capsule: `{handle}` contains `..` traversal — rejected"
2742 )
2743 .into());
2744 }
2745 std::path::Component::RootDir | std::path::Component::Prefix(_) => {
2746 return Err(format!(
2747 "expand_capsule: `{handle}` is an absolute path — only repo-relative paths are supported"
2748 )
2749 .into());
2750 }
2751 _ => {}
2752 }
2753 }
2754 let full_path = repo_root.join(path);
2755 let bytes = std::fs::read(&full_path)
2756 .map_err(|e| format!("expand_capsule: could not read `{rel_path}`: {e}"))?;
2757 let bounded = if bytes.len() > FILE_EXPAND_CAP_BYTES {
2759 let mut end = FILE_EXPAND_CAP_BYTES;
2760 while end > 0 && (bytes[end] & 0xC0) == 0x80 {
2762 end -= 1;
2763 }
2764 let s = String::from_utf8_lossy(&bytes[..end]);
2765 format!(
2766 "{s}\n[... truncated at {FILE_EXPAND_CAP_BYTES} bytes; call expand_capsule again with a line range if needed]"
2767 )
2768 } else {
2769 String::from_utf8_lossy(&bytes).into_owned()
2770 };
2771 Ok(bounded)
2772 } else if handle.starts_with("run:") {
2773 Err(format!(
2774 "expand_capsule: `run:` handle expansion is not yet supported (handle: `{handle}`)"
2775 )
2776 .into())
2777 } else {
2778 Err(format!(
2779 "expand_capsule: unrecognised handle format `{handle}`; \
2780 expected `memory:<id>`, `file:<path>`, or `run:<id>`"
2781 )
2782 .into())
2783 }
2784}
2785
2786pub fn rerank_capsules(
2795 query: &str,
2796 capsules: Vec<ContextCapsule>,
2797 reranker: &dyn crate::embeddings::Reranker,
2798 floor: f32,
2799 cap: usize,
2800) -> Vec<ContextCapsule> {
2801 if capsules.is_empty() {
2802 return capsules;
2803 }
2804
2805 let docs: Vec<&str> = capsules.iter().map(|c| c.summary.as_str()).collect();
2811 let scores = match reranker.rerank(query, &docs) {
2812 Ok(s) if s.len() == docs.len() => s,
2817 _ => {
2818 let mut out = capsules;
2820 if cap > 0 && out.len() > cap {
2821 out.truncate(cap);
2822 }
2823 return out;
2824 }
2825 };
2826
2827 let mut ranked: Vec<ContextCapsule> = capsules
2828 .into_iter()
2829 .zip(scores)
2830 .map(|(mut c, s)| {
2831 c.score = s;
2832 c
2833 })
2834 .collect();
2835
2836 ranked.sort_by(|a, b| {
2837 b.score
2838 .partial_cmp(&a.score)
2839 .unwrap_or(std::cmp::Ordering::Equal)
2840 });
2841
2842 ranked.retain(|c| c.score >= floor);
2843
2844 if cap > 0 && ranked.len() > cap {
2845 ranked.truncate(cap);
2846 }
2847
2848 ranked
2849}
2850
2851#[cfg(test)]
2852mod tests {
2853 use super::*;
2854
2855 fn capsule(kind: &str, summary: &str) -> ContextCapsule {
2856 ContextCapsule {
2857 id: "c".into(),
2858 kind: kind.into(),
2859 summary: summary.into(),
2860 token_estimate: 1,
2861 expansion_handle: "memory:x".into(),
2862 provenance: vec![],
2863 confidence: 1.0,
2864 freshness: 1.0,
2865 relevance: 1.0,
2866 scope_weight: 1.0,
2867 score: 1.0,
2868 }
2869 }
2870
2871 fn make_test_dir(tag: &str) -> std::path::PathBuf {
2874 use std::time::{SystemTime, UNIX_EPOCH};
2875 let ts = SystemTime::now()
2876 .duration_since(UNIX_EPOCH)
2877 .map(|d| d.subsec_nanos())
2878 .unwrap_or(0);
2879 let dir = std::env::temp_dir().join(format!("kbrain_test_{tag}_{ts}"));
2880 std::fs::create_dir_all(&dir).expect("create test dir");
2881 dir
2882 }
2883
2884 #[test]
2885 fn capsule_matches_kind_reads_memory_summary_prefix() {
2886 let mem = capsule("memory", "project:failure_pattern - linker not found");
2888 assert!(capsule_matches_kind(&mem, "failure_pattern"));
2889 assert!(!capsule_matches_kind(&mem, "command"));
2890 let repo = capsule("repo_file", "src/lib.rs:command - run build");
2892 assert!(capsule_matches_kind(&repo, "repo_file"));
2893 assert!(!capsule_matches_kind(&repo, "command"));
2894 }
2895
2896 #[test]
2899 fn usefulness_multiplier_neutral_at_zero_uses() {
2900 assert!((usefulness_multiplier(0.0, 0) - 1.0).abs() < f32::EPSILON);
2902 assert!((usefulness_multiplier(5.0, 0) - 1.0).abs() < f32::EPSILON);
2903 assert!((usefulness_multiplier(-5.0, 0) - 1.0).abs() < f32::EPSILON);
2904 }
2905
2906 #[test]
2910 fn usefulness_multiplier_blends_smoothly_in_transition() {
2911 let one_use = usefulness_multiplier(1.0, 1);
2914 assert!((one_use - 1.166_666_6).abs() < 1e-4, "got {one_use}");
2915 let two_uses = usefulness_multiplier(2.0, 2);
2918 assert!((two_uses - 1.333_333_4).abs() < 1e-4, "got {two_uses}");
2919 let two_uses_bad = usefulness_multiplier(-2.0, 2);
2921 assert!(
2923 (two_uses_bad - 0.666_666_7).abs() < 1e-4,
2924 "got {two_uses_bad}"
2925 );
2926 }
2927
2928 #[test]
2932 fn usefulness_multiplier_maps_ratio_onto_envelope() {
2933 assert!((usefulness_multiplier(5.0, 5) - 1.5).abs() < f32::EPSILON);
2935 assert!((usefulness_multiplier(-5.0, 5) - 0.5).abs() < f32::EPSILON);
2937 let mid = usefulness_multiplier(0.0, 6);
2939 assert!((mid - 1.0).abs() < f32::EPSILON, "got {mid}");
2940 let high = usefulness_multiplier(2.0, 4);
2942 assert!((high - 1.25).abs() < f32::EPSILON, "got {high}");
2943 let low = usefulness_multiplier(-2.0, 4);
2945 assert!((low - 0.75).abs() < f32::EPSILON, "got {low}");
2946 }
2947
2948 #[test]
2952 fn usefulness_multiplier_clamps_to_envelope() {
2953 assert!((usefulness_multiplier(100.0, 5) - 1.5).abs() < f32::EPSILON);
2955 assert!((usefulness_multiplier(-100.0, 5) - 0.5).abs() < f32::EPSILON);
2957 }
2958
2959 #[test]
2966 fn boost_gain_is_capped_so_cited_junk_cannot_beat_relevant_uncited() {
2967 let junk = apply_usefulness_boost(0.39, 1.5);
2968 let true_match = apply_usefulness_boost(0.53, 1.0);
2969 assert!(
2970 junk < true_match,
2971 "capped boost must preserve relevance order: junk {junk} vs match {true_match}"
2972 );
2973 assert!(junk <= 0.39 + USEFULNESS_BOOST_CAP + f32::EPSILON);
2975 }
2976
2977 #[test]
2981 fn boost_still_reorders_within_a_relevance_band() {
2982 let proven = apply_usefulness_boost(0.85, 1.5);
2983 let neutral = apply_usefulness_boost(0.90, 1.0);
2984 assert!(
2985 proven > neutral,
2986 "capped boost must still reorder near-equals: proven {proven} vs neutral {neutral}"
2987 );
2988 }
2989
2990 #[test]
2994 fn penalty_side_remains_multiplicative() {
2995 let penalized = apply_usefulness_boost(0.8, 0.5);
2996 assert!((penalized - 0.4).abs() < 1e-6);
2997 }
2998
2999 #[test]
3002 fn query_tokens_expands_build_class() {
3003 let toks = query_tokens("Build the project from source");
3004 assert!(toks.iter().any(|t| t == "build"));
3005 assert!(toks.iter().any(|t| t == "shell_background"));
3007 assert!(toks.iter().any(|t| t == "long_running"));
3008 }
3009
3010 #[test]
3011 fn query_tokens_expands_edit_class() {
3012 let toks = query_tokens("Modify the config to fix the bug");
3013 assert!(toks.iter().any(|t| t == "edit_file"));
3014 assert!(toks.iter().any(|t| t == "apply_patch"));
3015 }
3016
3017 #[test]
3018 fn query_tokens_expands_search_class() {
3019 let toks = query_tokens("Find all references to the symbol");
3020 assert!(toks.iter().any(|t| t == "glob"));
3021 assert!(toks.iter().any(|t| t == "search_files"));
3022 }
3023
3024 #[test]
3025 fn query_tokens_no_expansion_on_unrelated_query() {
3026 let toks = query_tokens("hello world testing nothing");
3027 assert!(toks.iter().any(|t| t == "hello"));
3029 assert!(toks.iter().any(|t| t == "world"));
3031 }
3032
3033 #[test]
3036 fn jaccard_is_zero_for_disjoint_sets() {
3037 let a: std::collections::HashSet<String> =
3038 ["foo", "bar"].iter().map(|s| s.to_string()).collect();
3039 let b: std::collections::HashSet<String> =
3040 ["baz", "qux"].iter().map(|s| s.to_string()).collect();
3041 assert!((jaccard(&a, &b) - 0.0).abs() < f32::EPSILON);
3042 }
3043
3044 #[test]
3045 fn jaccard_is_one_for_identical_sets() {
3046 let a: std::collections::HashSet<String> =
3047 ["foo", "bar"].iter().map(|s| s.to_string()).collect();
3048 let b = a.clone();
3049 assert!((jaccard(&a, &b) - 1.0).abs() < f32::EPSILON);
3050 }
3051
3052 #[test]
3053 fn jaccard_partial_overlap() {
3054 let a: std::collections::HashSet<String> = ["foo", "bar", "baz"]
3055 .iter()
3056 .map(|s| s.to_string())
3057 .collect();
3058 let b: std::collections::HashSet<String> =
3059 ["bar", "qux"].iter().map(|s| s.to_string()).collect();
3060 assert!((jaccard(&a, &b) - 0.25).abs() < f32::EPSILON);
3062 }
3063
3064 #[test]
3065 fn summary_token_set_lowercases_and_filters_short() {
3066 let set = summary_token_set("Build the Foo-bar project");
3067 assert!(set.contains("build"));
3068 assert!(set.contains("foo"));
3069 assert!(set.contains("bar"));
3070 assert!(set.contains("project"));
3071 assert!(set.contains("the"));
3073 }
3074
3075 fn insert_memory_with_embedding(
3081 conn: &rusqlite::Connection,
3082 memory_id: &str,
3083 text: &str,
3084 embedder: &dyn embeddings::Embedder,
3085 ) {
3086 let normalized = kimetsu_core::memory::normalize_memory_text(text);
3087 conn.execute(
3088 "
3089 INSERT INTO memories (
3090 memory_id, scope, kind, text, normalized_text, confidence,
3091 source_event_id, provenance_snapshot_json, created_at,
3092 use_count, usefulness_score, embedding, embedding_model
3093 )
3094 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
3095 '2026-05-01T00:00:00Z', 0, 0.0, ?4, ?5)
3096 ",
3097 rusqlite::params![
3098 memory_id,
3099 text,
3100 normalized,
3101 embeddings::encode_embedding(&embedder.embed(text).expect("embed test row")),
3102 embedder.model_id(),
3103 ],
3104 )
3105 .expect("insert memory");
3106 conn.execute(
3107 "INSERT INTO memories_fts (memory_id, text, kind, scope) VALUES (?1, ?2, 'fact', 'global_user')",
3108 rusqlite::params![memory_id, text],
3109 )
3110 .expect("insert fts row");
3111 }
3112
3113 #[test]
3123 fn hybrid_retrieval_uses_cosine_score_to_rerank() {
3124 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3125 crate::schema::initialize(&conn).expect("init schema");
3126 let stub = embeddings::StubEmbedder::new();
3127
3128 insert_memory_with_embedding(&conn, "m_rg", "use ripgrep for code search", &stub);
3129 insert_memory_with_embedding(
3130 &conn,
3131 "m_unrelated",
3132 "cookie recipe with chocolate chips",
3133 &stub,
3134 );
3135
3136 let weights = kimetsu_core::config::BrokerWeights::default();
3139 let bundle = retrieve_context_with_embedder(
3140 &conn,
3141 "/fake-repo",
3142 &weights,
3143 ContextRequest {
3144 stage: "localization".to_string(),
3145 query: "ripgrep search".to_string(),
3146 budget_tokens: 4000,
3147 ..Default::default()
3148 },
3149 &[],
3150 &stub,
3151 )
3152 .expect("retrieve");
3153
3154 let memory_handles: Vec<_> = bundle
3155 .capsules
3156 .iter()
3157 .filter(|c| c.expansion_handle.starts_with("memory:"))
3158 .collect();
3159 assert!(
3160 !memory_handles.is_empty(),
3161 "at least one memory should surface"
3162 );
3163 assert_eq!(
3165 memory_handles[0].expansion_handle,
3166 "memory:m_rg",
3167 "ripgrep memory should outrank the cookie recipe; ranked: {:?}",
3168 memory_handles
3169 .iter()
3170 .map(|c| &c.expansion_handle)
3171 .collect::<Vec<_>>()
3172 );
3173 }
3174
3175 #[test]
3182 fn hybrid_retrieval_skips_cosine_on_model_id_mismatch() {
3183 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3184 crate::schema::initialize(&conn).expect("init schema");
3185 let stub = embeddings::StubEmbedder::new();
3186 insert_memory_with_embedding(&conn, "m_xref", "use ripgrep for code search", &stub);
3187
3188 conn.execute(
3193 "UPDATE memories SET embedding_model = 'bge-small-en-v1.5' WHERE memory_id = 'm_xref'",
3194 [],
3195 )
3196 .expect("force model_id mismatch");
3197
3198 let weights = kimetsu_core::config::BrokerWeights::default();
3203 let bundle = retrieve_context_with_embedder(
3204 &conn,
3205 "/fake-repo",
3206 &weights,
3207 ContextRequest {
3208 stage: "localization".to_string(),
3209 query: "ripgrep search".to_string(),
3210 budget_tokens: 4000,
3211 ..Default::default()
3212 },
3213 &[],
3214 &stub,
3215 )
3216 .expect("retrieve");
3217
3218 assert!(
3219 bundle
3220 .capsules
3221 .iter()
3222 .any(|c| c.expansion_handle == "memory:m_xref"),
3223 "cross-model row should still match lexically (cosine skipped, FTS works)"
3224 );
3225 }
3226
3227 #[test]
3234 fn usefulness_decay_disabled_when_half_life_is_zero_or_negative() {
3235 let ancient = "2021-01-01T00:00:00Z";
3237 assert!((usefulness_decay(Some(ancient), ancient, 0.0) - 1.0).abs() < f32::EPSILON);
3238 assert!((usefulness_decay(Some(ancient), ancient, -1.0) - 1.0).abs() < f32::EPSILON);
3239 }
3240
3241 #[test]
3245 fn usefulness_decay_returns_one_on_unparseable_timestamps() {
3246 assert!(
3247 (usefulness_decay(Some("not-a-date"), "also-not", 30.0) - 1.0).abs() < f32::EPSILON
3248 );
3249 }
3250
3251 #[test]
3254 fn usefulness_decay_full_at_zero_age() {
3255 let future = "2099-01-01T00:00:00Z";
3257 let d = usefulness_decay(Some(future), future, 30.0);
3258 assert!((d - 1.0).abs() < f32::EPSILON, "got {d}");
3259 }
3260
3261 #[test]
3266 fn usefulness_decay_follows_half_life_curve() {
3267 let half_life = 10.0_f32;
3268 let now = OffsetDateTime::now_utc();
3269 let fmt = &time::format_description::well_known::Rfc3339;
3270
3271 let one_half_life_ago = (now - time::Duration::seconds((half_life * 86_400.0) as i64))
3273 .format(fmt)
3274 .expect("format");
3275 let d1 = usefulness_decay(Some(&one_half_life_ago), &one_half_life_ago, half_life);
3276 assert!(
3277 (d1 - 0.5).abs() < 0.01,
3278 "expected ~0.5 at one half-life, got {d1}"
3279 );
3280
3281 let two_half_lives_ago = (now
3283 - time::Duration::seconds((2.0 * half_life * 86_400.0) as i64))
3284 .format(fmt)
3285 .expect("format");
3286 let d2 = usefulness_decay(Some(&two_half_lives_ago), &two_half_lives_ago, half_life);
3287 assert!(
3288 (d2 - 0.25).abs() < 0.01,
3289 "expected ~0.25 at two half-lives, got {d2}"
3290 );
3291 }
3292
3293 #[test]
3297 fn usefulness_decay_falls_back_to_created_at_when_last_useful_is_none() {
3298 let now = OffsetDateTime::now_utc();
3299 let fmt = &time::format_description::well_known::Rfc3339;
3300 let one_day_ago = (now - time::Duration::seconds(86_400))
3301 .format(fmt)
3302 .expect("format");
3303 let d = usefulness_decay(None, &one_day_ago, 30.0);
3304 assert!(
3306 (d - 0.977).abs() < 0.01,
3307 "expected ~0.977 for 1-day-old created_at under 30d half-life, got {d}"
3308 );
3309 }
3310
3311 #[test]
3316 fn aged_cited_memory_ranks_below_recently_cited_memory() {
3317 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3318 crate::schema::initialize(&conn).expect("init schema");
3319
3320 let now = OffsetDateTime::now_utc();
3321 let fmt = &time::format_description::well_known::Rfc3339;
3322 let one_day_ago = (now - time::Duration::seconds(86_400))
3323 .format(fmt)
3324 .expect("format");
3325 let one_year_ago = (now - time::Duration::seconds(365 * 86_400))
3326 .format(fmt)
3327 .expect("format");
3328
3329 for (mid, last_useful) in [("m_recent", &one_day_ago), ("m_aged", &one_year_ago)] {
3333 let text = "use ripgrep for code search";
3334 let normalized = kimetsu_core::memory::normalize_memory_text(text);
3335 conn.execute(
3336 "
3337 INSERT INTO memories (
3338 memory_id, scope, kind, text, normalized_text, confidence,
3339 source_event_id, provenance_snapshot_json, created_at,
3340 use_count, usefulness_score, last_useful_at
3341 )
3342 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
3343 '2024-01-01T00:00:00Z', 5, 5.0, ?4)
3344 ",
3345 rusqlite::params![mid, text, normalized, last_useful],
3346 )
3347 .expect("insert memory");
3348 conn.execute(
3349 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3350 VALUES (?1, ?2, 'fact', 'global_user')",
3351 rusqlite::params![mid, text],
3352 )
3353 .expect("insert fts");
3354 }
3355
3356 let weights = kimetsu_core::config::BrokerWeights::default();
3358 let bundle = retrieve_context_with_embedder(
3359 &conn,
3360 "/fake-repo",
3361 &weights,
3362 ContextRequest {
3363 stage: "localization".to_string(),
3364 query: "ripgrep search".to_string(),
3365 budget_tokens: 4000,
3366 ..Default::default()
3367 },
3368 &[],
3369 &embeddings::NoopEmbedder,
3370 )
3371 .expect("retrieve");
3372
3373 let mem_order: Vec<&str> = bundle
3374 .capsules
3375 .iter()
3376 .filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
3377 .collect();
3378 assert_eq!(
3379 mem_order.first().copied(),
3380 Some("m_recent"),
3381 "recently-cited memory must rank first under decay; got order {mem_order:?}"
3382 );
3383 }
3384
3385 #[test]
3390 fn aged_cited_memory_does_not_decay_when_half_life_is_zero() {
3391 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3392 crate::schema::initialize(&conn).expect("init schema");
3393
3394 let now = OffsetDateTime::now_utc();
3395 let fmt = &time::format_description::well_known::Rfc3339;
3396 let one_day_ago = (now - time::Duration::seconds(86_400))
3397 .format(fmt)
3398 .expect("format");
3399 let one_year_ago = (now - time::Duration::seconds(365 * 86_400))
3400 .format(fmt)
3401 .expect("format");
3402
3403 for (mid, last_useful) in [("m_recent", &one_day_ago), ("m_aged", &one_year_ago)] {
3404 let text = "use ripgrep for code search";
3405 let normalized = kimetsu_core::memory::normalize_memory_text(text);
3406 conn.execute(
3407 "
3408 INSERT INTO memories (
3409 memory_id, scope, kind, text, normalized_text, confidence,
3410 source_event_id, provenance_snapshot_json, created_at,
3411 use_count, usefulness_score, last_useful_at
3412 )
3413 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
3414 '2024-01-01T00:00:00Z', 5, 5.0, ?4)
3415 ",
3416 rusqlite::params![mid, text, normalized, last_useful],
3417 )
3418 .expect("insert memory");
3419 conn.execute(
3420 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3421 VALUES (?1, ?2, 'fact', 'global_user')",
3422 rusqlite::params![mid, text],
3423 )
3424 .expect("insert fts");
3425 }
3426
3427 let weights = kimetsu_core::config::BrokerWeights {
3429 decay_half_life_days: 0.0,
3430 ..Default::default()
3431 };
3432
3433 let bundle = retrieve_context_with_embedder(
3434 &conn,
3435 "/fake-repo",
3436 &weights,
3437 ContextRequest {
3438 stage: "localization".to_string(),
3439 query: "ripgrep search".to_string(),
3440 budget_tokens: 4000,
3441 ..Default::default()
3442 },
3443 &[],
3444 &embeddings::NoopEmbedder,
3445 )
3446 .expect("retrieve");
3447
3448 let scores: Vec<(String, f32)> = bundle
3454 .capsules
3455 .iter()
3456 .filter_map(|c| {
3457 c.expansion_handle
3458 .strip_prefix("memory:")
3459 .map(|id| (id.to_string(), c.score))
3460 })
3461 .collect();
3462 assert_eq!(scores.len(), 2, "both memories should surface");
3463 let recent_score = scores
3464 .iter()
3465 .find(|(id, _)| id == "m_recent")
3466 .map(|(_, s)| *s)
3467 .expect("m_recent present");
3468 let aged_score = scores
3469 .iter()
3470 .find(|(id, _)| id == "m_aged")
3471 .map(|(_, s)| *s)
3472 .expect("m_aged present");
3473 assert!(
3475 (recent_score - aged_score).abs() < 1e-4,
3476 "with decay disabled the two memories should tie on score: recent={recent_score} aged={aged_score}"
3477 );
3478 }
3479
3480 #[test]
3485 fn hybrid_retrieval_with_noop_embedder_is_lexical_only() {
3486 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3487 crate::schema::initialize(&conn).expect("init schema");
3488 let stub = embeddings::StubEmbedder::new();
3489 insert_memory_with_embedding(&conn, "m_a", "use ripgrep", &stub);
3491 insert_memory_with_embedding(&conn, "m_b", "use ripgrep too", &stub);
3492
3493 let weights = kimetsu_core::config::BrokerWeights::default();
3496 let bundle = retrieve_context_with_embedder(
3497 &conn,
3498 "/fake-repo",
3499 &weights,
3500 ContextRequest {
3501 stage: "localization".to_string(),
3502 query: "ripgrep".to_string(),
3503 budget_tokens: 4000,
3504 ..Default::default()
3505 },
3506 &[],
3507 &embeddings::NoopEmbedder,
3508 )
3509 .expect("retrieve");
3510
3511 let count = bundle
3512 .capsules
3513 .iter()
3514 .filter(|c| c.expansion_handle.starts_with("memory:"))
3515 .count();
3516 assert_eq!(count, 2, "both memories should surface via FTS");
3517 }
3518
3519 #[cfg(feature = "embeddings")]
3546 #[test]
3547 fn ann_finds_semantic_match_fts_misses() {
3548 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3549 crate::schema::initialize(&conn).expect("init schema");
3550
3551 struct OracleEmbedder;
3554 impl embeddings::Embedder for OracleEmbedder {
3555 fn embed(&self, _text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
3556 Ok(vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
3558 }
3559 fn model_id(&self) -> &str {
3560 "oracle-d8"
3561 }
3562 fn dim(&self) -> usize {
3563 8
3564 }
3565 }
3566
3567 let model_id = "oracle-d8";
3568
3569 let sem_vec = embeddings::encode_embedding(&[1.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
3572 let sem_text = "cookie recipe chocolate";
3573 let sem_norm = kimetsu_core::memory::normalize_memory_text(sem_text);
3574 conn.execute(
3575 "INSERT INTO memories (
3576 memory_id, scope, kind, text, normalized_text, confidence,
3577 source_event_id, provenance_snapshot_json, created_at,
3578 use_count, usefulness_score, embedding, embedding_model
3579 )
3580 VALUES ('m_semantic', 'global_user', 'fact', ?1, ?2, 1.0, NULL, '{}',
3581 '2026-01-01T00:00:00Z', 0, 0.0, ?3, ?4)",
3582 rusqlite::params![sem_text, sem_norm, sem_vec, model_id],
3583 )
3584 .expect("insert m_semantic");
3585 conn.execute(
3586 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3587 VALUES ('m_semantic', ?1, 'fact', 'global_user')",
3588 rusqlite::params![sem_text],
3589 )
3590 .expect("insert m_semantic fts");
3591
3592 let decoy_vec = embeddings::encode_embedding(&[0.0f32, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
3594 let decoy_text = "git rebase squash commits";
3595 let decoy_norm = kimetsu_core::memory::normalize_memory_text(decoy_text);
3596 conn.execute(
3597 "INSERT INTO memories (
3598 memory_id, scope, kind, text, normalized_text, confidence,
3599 source_event_id, provenance_snapshot_json, created_at,
3600 use_count, usefulness_score, embedding, embedding_model
3601 )
3602 VALUES ('m_decoy', 'global_user', 'fact', ?1, ?2, 1.0, NULL, '{}',
3603 '2026-01-01T00:00:00Z', 0, 0.0, ?3, ?4)",
3604 rusqlite::params![decoy_text, decoy_norm, decoy_vec, model_id],
3605 )
3606 .expect("insert m_decoy");
3607 conn.execute(
3608 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3609 VALUES ('m_decoy', ?1, 'fact', 'global_user')",
3610 rusqlite::params![decoy_text],
3611 )
3612 .expect("insert m_decoy fts");
3613
3614 let fts_hits: i64 = conn
3616 .query_row(
3617 "SELECT COUNT(*) FROM memories_fts \
3618 WHERE memories_fts MATCH 'phosphorescent bioluminescent'",
3619 [],
3620 |r| r.get(0),
3621 )
3622 .unwrap_or(0);
3623 assert_eq!(
3624 fts_hits, 0,
3625 "sanity: query tokens must not appear in any memory text"
3626 );
3627
3628 let weights = kimetsu_core::config::BrokerWeights::default();
3632 let bundle = retrieve_context_with_embedder(
3633 &conn,
3634 "/fake-repo",
3635 &weights,
3636 ContextRequest {
3637 stage: "localization".to_string(),
3638 query: "phosphorescent bioluminescent organism".to_string(),
3639 budget_tokens: 4000,
3640 ..Default::default()
3641 },
3642 &[],
3643 &OracleEmbedder,
3644 )
3645 .expect("retrieve");
3646
3647 let handles: Vec<&str> = bundle
3648 .capsules
3649 .iter()
3650 .filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
3651 .collect();
3652
3653 assert!(
3654 handles.contains(&"m_semantic"),
3655 "ANN must surface m_semantic (cosine=1 with oracle query) even though \
3656 FTS found nothing; got handles: {handles:?}"
3657 );
3658 }
3659
3660 #[cfg(feature = "embeddings")]
3662 #[test]
3663 fn dedup_memory_matched_by_fts_and_ann_appears_once() {
3664 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3665 crate::schema::initialize(&conn).expect("init schema");
3666
3667 let stub = embeddings::StubEmbedder::new();
3668
3669 insert_memory_with_embedding(&conn, "m_both", "use ripgrep for fast search", &stub);
3673
3674 let weights = kimetsu_core::config::BrokerWeights::default();
3675 let bundle = retrieve_context_with_embedder(
3676 &conn,
3677 "/fake-repo",
3678 &weights,
3679 ContextRequest {
3680 stage: "localization".to_string(),
3681 query: "ripgrep".to_string(),
3682 budget_tokens: 4000,
3683 ..Default::default()
3684 },
3685 &[],
3686 &stub,
3687 )
3688 .expect("retrieve");
3689
3690 let count = bundle
3691 .capsules
3692 .iter()
3693 .filter(|c| c.expansion_handle == "memory:m_both")
3694 .count();
3695 assert_eq!(
3696 count,
3697 1,
3698 "m_both (matched by both FTS and ANN) must appear exactly once; \
3699 bundle: {:?}",
3700 bundle
3701 .capsules
3702 .iter()
3703 .map(|c| &c.expansion_handle)
3704 .collect::<Vec<_>>()
3705 );
3706 }
3707
3708 #[cfg(feature = "embeddings")]
3732 #[test]
3733 fn embedding_mmr_collapses_paraphrases_but_jaccard_does_not() {
3734 struct OracleEmbedder;
3737 impl embeddings::Embedder for OracleEmbedder {
3738 fn embed(&self, _text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
3739 let mut v = vec![0.0f32; 8];
3740 v[0] = 1.0;
3741 Ok(v)
3742 }
3743 fn model_id(&self) -> &str {
3744 "oracle-d8"
3745 }
3746 fn dim(&self) -> usize {
3747 8
3748 }
3749 }
3750
3751 let oracle = OracleEmbedder;
3754 let weights = kimetsu_core::config::BrokerWeights::default();
3755
3756 let m_rg1_text = "prefer ripgrep for searching source code";
3759 let m_rg2_text = "rg is the fastest way to locate patterns";
3760
3761 let conn = rusqlite::Connection::open_in_memory().expect("in-memory");
3767 crate::schema::initialize(&conn).expect("init schema");
3768 insert_memory_with_embedding(&conn, "m_rg1", m_rg1_text, &oracle);
3769 insert_memory_with_embedding(&conn, "m_rg2", m_rg2_text, &oracle);
3770
3771 let bundle_embedding = retrieve_context_with_embedder(
3772 &conn,
3773 "/fake-repo",
3774 &weights,
3775 ContextRequest {
3776 stage: "localization".to_string(),
3777 query: "search source patterns".to_string(),
3779 budget_tokens: 20_000,
3780 max_capsules: 1, ..Default::default()
3782 },
3783 &[],
3784 &oracle,
3785 )
3786 .expect("retrieve with oracle embedder");
3787
3788 let emb_in_capsules = bundle_embedding
3791 .capsules
3792 .iter()
3793 .filter(|c| {
3794 c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
3795 })
3796 .count();
3797 assert_eq!(
3798 emb_in_capsules,
3799 1,
3800 "embedding-MMR must collapse cosine=1.0 paraphrases: with max_capsules=1 \
3801 only ONE should be included; capsule handles: {:?}; excluded: {:?}",
3802 bundle_embedding
3803 .capsules
3804 .iter()
3805 .map(|c| &c.expansion_handle)
3806 .collect::<Vec<_>>(),
3807 bundle_embedding
3808 .excluded
3809 .iter()
3810 .map(|c| &c.expansion_handle)
3811 .collect::<Vec<_>>()
3812 );
3813
3814 let emb_in_excluded = bundle_embedding
3816 .excluded
3817 .iter()
3818 .filter(|c| {
3819 c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
3820 })
3821 .count();
3822 assert_eq!(
3823 emb_in_excluded,
3824 1,
3825 "the second near-duplicate must be in excluded under embedding-MMR; \
3826 excluded handles: {:?}",
3827 bundle_embedding
3828 .excluded
3829 .iter()
3830 .map(|c| &c.expansion_handle)
3831 .collect::<Vec<_>>()
3832 );
3833
3834 let conn2 = rusqlite::Connection::open_in_memory().expect("in-memory 2");
3839 crate::schema::initialize(&conn2).expect("init schema 2");
3840 insert_memory_with_embedding(&conn2, "m_rg1", m_rg1_text, &oracle);
3841 insert_memory_with_embedding(&conn2, "m_rg2", m_rg2_text, &oracle);
3842
3843 let bundle_lean = retrieve_context_with_embedder(
3844 &conn2,
3845 "/fake-repo",
3846 &weights,
3847 ContextRequest {
3848 stage: "localization".to_string(),
3849 query: "search source patterns".to_string(),
3850 budget_tokens: 20_000,
3851 max_capsules: 2, ..Default::default()
3853 },
3854 &[],
3855 &embeddings::NoopEmbedder,
3856 )
3857 .expect("retrieve with NoopEmbedder");
3858
3859 let lean_in_capsules = bundle_lean
3860 .capsules
3861 .iter()
3862 .filter(|c| {
3863 c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
3864 })
3865 .count();
3866 assert_eq!(
3867 lean_in_capsules,
3868 2,
3869 "Jaccard-only path must NOT collapse the two paraphrases (different words, \
3870 low token overlap → both survive MMR with max_capsules=2); capsule handles: {:?}",
3871 bundle_lean
3872 .capsules
3873 .iter()
3874 .map(|c| &c.expansion_handle)
3875 .collect::<Vec<_>>()
3876 );
3877 }
3878
3879 #[test]
3882 fn content_tokens_strips_stopwords_keeps_topical_words() {
3883 let got = content_tokens("Tell me about kimetsu, what's the idea of the repo");
3884 assert_eq!(got, vec!["kimetsu", "idea", "repo"]);
3887 }
3888
3889 #[test]
3890 fn light_stem_strips_one_inflection_suffix() {
3891 assert_eq!(light_stem("benchmarked"), "benchmark");
3892 assert_eq!(light_stem("benchmarking"), "benchmark");
3893 assert_eq!(light_stem("repos"), "repo");
3894 assert_eq!(light_stem("does"), "does");
3896 assert_eq!(light_stem("toml"), "toml");
3897 }
3898
3899 #[test]
3906 fn stemmed_query_matches_inflected_corpus_through_floor() {
3907 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3908 crate::schema::initialize(&conn).expect("init schema");
3909 let insert = |id: &str, text: &str| {
3910 let norm = kimetsu_core::memory::normalize_memory_text(text);
3911 conn.execute(
3912 "INSERT INTO memories (
3913 memory_id, scope, kind, text, normalized_text, confidence,
3914 source_event_id, provenance_snapshot_json, created_at,
3915 use_count, usefulness_score, embedding, embedding_model
3916 )
3917 VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
3918 '2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
3919 rusqlite::params![id, text, norm],
3920 )
3921 .expect("insert memory");
3922 conn.execute(
3923 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3924 VALUES (?1, ?2, 'fact', 'global_user')",
3925 rusqlite::params![id, text],
3926 )
3927 .expect("insert fts");
3928 };
3929 insert(
3930 "m_bench",
3931 "kimetsu benchmark runs go through the kbench binary and the Terminal-Bench driver",
3932 );
3933 insert(
3934 "m_doctor",
3935 "kimetsu doctor version-skew check parses process start times on Windows via CIM",
3936 );
3937 insert(
3938 "m_gc",
3939 "kimetsu runs auto-GC on run creation; keep the env guard at the trigger site",
3940 );
3941
3942 let bundle = retrieve_context_with_embedder(
3943 &conn,
3944 "/fake-repo",
3945 &kimetsu_core::config::BrokerWeights::default(),
3946 ContextRequest {
3947 stage: "localization".to_string(),
3948 query: "Can you find out how kimetsu is benchmarked?".to_string(),
3949 budget_tokens: 2000,
3950 max_capsules: 2,
3951 min_lexical_coverage: 0.5,
3952 ..Default::default()
3953 },
3954 &[],
3955 &embeddings::NoopEmbedder,
3956 )
3957 .expect("retrieve");
3958 let handles: Vec<_> = bundle
3959 .capsules
3960 .iter()
3961 .map(|c| c.expansion_handle.as_str())
3962 .collect();
3963 assert!(
3964 handles.contains(&"memory:m_bench"),
3965 "stemmed 'benchmarked' must surface the benchmark memory; got {handles:?}"
3966 );
3967 assert!(
3968 !handles.contains(&"memory:m_doctor") && !handles.contains(&"memory:m_gc"),
3969 "off-topic memories sharing only 'kimetsu' must stay below the floor; got {handles:?}"
3970 );
3971 }
3972
3973 #[test]
3974 fn weighted_coverage_ignores_zero_idf_tokens() {
3975 let content = vec![
3979 "kimetsu".to_string(),
3980 "idea".to_string(),
3981 "repo".to_string(),
3982 ];
3983 let mut idf = HashMap::new();
3984 idf.insert("kimetsu".to_string(), 0.0);
3985 idf.insert("idea".to_string(), 1.386);
3986 idf.insert("repo".to_string(), 0.693);
3987
3988 let cov = weighted_coverage(
3990 &content,
3991 &idf,
3992 "global:fact - the git repo and kimetsu brain",
3993 );
3994 assert!((cov - 0.333).abs() < 0.01, "got {cov}");
3995
3996 let cov_topical =
3998 weighted_coverage(&content, &idf, "global:fact - the core idea of kimetsu");
3999 assert!(cov_topical > 0.6, "got {cov_topical}");
4000 }
4001
4002 #[test]
4003 fn escape_like_neutralizes_wildcards() {
4004 assert_eq!(escape_like("a_b%c"), "a\\_b\\%c");
4005 assert_eq!(escape_like("plain"), "plain");
4006 }
4007
4008 #[test]
4022 fn lexical_floor_drops_offtopic_memories_sharing_project_name() {
4023 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4024 crate::schema::initialize(&conn).expect("init schema");
4025
4026 let insert = |id: &str, text: &str| {
4027 let norm = kimetsu_core::memory::normalize_memory_text(text);
4028 conn.execute(
4029 "INSERT INTO memories (
4030 memory_id, scope, kind, text, normalized_text, confidence,
4031 source_event_id, provenance_snapshot_json, created_at,
4032 use_count, usefulness_score, embedding, embedding_model
4033 )
4034 VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
4035 '2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
4036 rusqlite::params![id, text, norm],
4037 )
4038 .expect("insert memory");
4039 conn.execute(
4040 "INSERT INTO memories_fts (memory_id, text, kind, scope)
4041 VALUES (?1, ?2, 'fact', 'global_user')",
4042 rusqlite::params![id, text],
4043 )
4044 .expect("insert fts");
4045 };
4046
4047 insert(
4050 "m1",
4051 "When implementing a setup command that calls init_project, tests must call \
4052 git_init_boundary before setup_cmd so ProjectPaths discover resolves to the temp \
4053 dir instead of climbing to the real parent git repo including the user brain at kimetsu",
4054 );
4055 insert(
4056 "m2",
4057 "A member crate with default embeddings silently turned embeddings on for the entire \
4058 cargo test workspace build graph because cargo unifies features; kimetsu-chat \
4059 retrieval tests failed",
4060 );
4061 insert(
4062 "m3",
4063 "In toml 0.9 use toml from_str to parse a TOML document into a Value not str parse; \
4064 implementing config get and set in kimetsu-cli",
4065 );
4066
4067 let query = "Tell me about kimetsu, what's the idea of the repo".to_string();
4068 let weights = kimetsu_core::config::BrokerWeights::default();
4069 let handles = |bundle: &ContextBundle| {
4070 bundle
4071 .capsules
4072 .iter()
4073 .map(|c| c.expansion_handle.clone())
4074 .collect::<Vec<_>>()
4075 };
4076
4077 let no_floor = retrieve_context_with_embedder(
4079 &conn,
4080 "/fake-repo",
4081 &weights,
4082 ContextRequest {
4083 stage: "localization".to_string(),
4084 query: query.clone(),
4085 budget_tokens: 2000,
4086 max_capsules: 8,
4087 min_lexical_coverage: 0.0,
4088 ..Default::default()
4089 },
4090 &[],
4091 &embeddings::NoopEmbedder,
4092 )
4093 .expect("retrieve without floor");
4094 let before = handles(&no_floor);
4095 assert!(
4096 before.contains(&"memory:m2".to_string()) && before.contains(&"memory:m3".to_string()),
4097 "sanity: without the floor the pure-project-name memories should surface; got {before:?}"
4098 );
4099
4100 let floored = retrieve_context_with_embedder(
4102 &conn,
4103 "/fake-repo",
4104 &weights,
4105 ContextRequest {
4106 stage: "localization".to_string(),
4107 query,
4108 budget_tokens: 2000,
4109 max_capsules: 8,
4110 min_lexical_coverage: 0.5,
4111 ..Default::default()
4112 },
4113 &[],
4114 &embeddings::NoopEmbedder,
4115 )
4116 .expect("retrieve with floor");
4117 let after = handles(&floored);
4118 assert!(
4119 !after.contains(&"memory:m2".to_string()) && !after.contains(&"memory:m3".to_string()),
4120 "the lexical floor must drop memories whose only match is the corpus-ubiquitous \
4121 project name; surviving: {after:?}"
4122 );
4123 }
4124
4125 #[test]
4128 fn lexical_floor_keeps_ontopic_memory() {
4129 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4130 crate::schema::initialize(&conn).expect("init schema");
4131
4132 let insert = |id: &str, text: &str| {
4133 let norm = kimetsu_core::memory::normalize_memory_text(text);
4134 conn.execute(
4135 "INSERT INTO memories (
4136 memory_id, scope, kind, text, normalized_text, confidence,
4137 source_event_id, provenance_snapshot_json, created_at,
4138 use_count, usefulness_score, embedding, embedding_model
4139 )
4140 VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
4141 '2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
4142 rusqlite::params![id, text, norm],
4143 )
4144 .expect("insert memory");
4145 conn.execute(
4146 "INSERT INTO memories_fts (memory_id, text, kind, scope)
4147 VALUES (?1, ?2, 'fact', 'global_user')",
4148 rusqlite::params![id, text],
4149 )
4150 .expect("insert fts");
4151 };
4152
4153 insert(
4155 "d1",
4156 "The distiller runs at session end and harvests durable lessons from the transcript",
4157 );
4158 insert(
4159 "n1",
4160 "Unrelated note about git rebase and squashing commits",
4161 );
4162
4163 let bundle = retrieve_context_with_embedder(
4164 &conn,
4165 "/fake-repo",
4166 &kimetsu_core::config::BrokerWeights::default(),
4167 ContextRequest {
4168 stage: "localization".to_string(),
4169 query: "how does the distiller work".to_string(),
4170 budget_tokens: 2000,
4171 min_lexical_coverage: 0.5,
4172 ..Default::default()
4173 },
4174 &[],
4175 &embeddings::NoopEmbedder,
4176 )
4177 .expect("retrieve");
4178
4179 assert!(
4180 bundle
4181 .capsules
4182 .iter()
4183 .any(|c| c.expansion_handle == "memory:d1"),
4184 "on-topic memory covering the rare query word must survive the floor; got: {:?}",
4185 bundle
4186 .capsules
4187 .iter()
4188 .map(|c| &c.expansion_handle)
4189 .collect::<Vec<_>>()
4190 );
4191 }
4192
4193 #[cfg(feature = "embeddings")]
4202 #[test]
4203 fn min_semantic_score_floor_drops_off_topic_queries() {
4204 struct DirectionalEmbedder {
4215 marker: &'static str,
4217 }
4218 impl embeddings::Embedder for DirectionalEmbedder {
4219 fn embed(&self, text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
4220 let mut v = vec![0.0f32; 8];
4221 if text.contains(self.marker) {
4222 v[0] = 1.0;
4223 } else {
4224 v[1] = 1.0;
4225 }
4226 Ok(v)
4227 }
4228 fn model_id(&self) -> &str {
4229 "directional-d8"
4230 }
4231 fn dim(&self) -> usize {
4232 8
4233 }
4234 }
4235
4236 let emb = DirectionalEmbedder { marker: "TOPIC_A" };
4237
4238 let conn = rusqlite::Connection::open_in_memory().expect("in-memory");
4239 crate::schema::initialize(&conn).expect("init schema");
4240
4241 insert_memory_with_embedding(&conn, "m_b", "cookie recipe chocolate baking TOPIC_B", &emb);
4243
4244 let weights = kimetsu_core::config::BrokerWeights::default();
4245
4246 let bundle_off = retrieve_context_with_embedder(
4248 &conn,
4249 "/fake-repo",
4250 &weights,
4251 ContextRequest {
4252 stage: "localization".to_string(),
4253 query: "TOPIC_A unrelated phosphorescent".to_string(),
4255 budget_tokens: 4000,
4256 min_semantic_score: 0.1, ..Default::default()
4258 },
4259 &[],
4260 &emb,
4261 )
4262 .expect("retrieve off-topic");
4263
4264 assert!(
4265 bundle_off.capsules.is_empty(),
4266 "off-topic query (cosine=0 < floor=0.1) must produce zero capsules; \
4267 got: {:?}",
4268 bundle_off
4269 .capsules
4270 .iter()
4271 .map(|c| &c.expansion_handle)
4272 .collect::<Vec<_>>()
4273 );
4274
4275 let conn2 = rusqlite::Connection::open_in_memory().expect("in-memory 2");
4278 crate::schema::initialize(&conn2).expect("init schema 2");
4279 insert_memory_with_embedding(
4280 &conn2,
4281 "m_b2",
4282 "cookie recipe chocolate TOPIC_B baking"
4283 .to_string()
4284 .as_str(),
4285 &emb,
4286 );
4287
4288 let bundle_on = retrieve_context_with_embedder(
4289 &conn2,
4290 "/fake-repo",
4291 &weights,
4292 ContextRequest {
4293 stage: "localization".to_string(),
4294 query: "cookie chocolate TOPIC_B".to_string(),
4296 budget_tokens: 4000,
4297 min_semantic_score: 0.1,
4298 ..Default::default()
4299 },
4300 &[],
4301 &emb,
4302 )
4303 .expect("retrieve on-topic");
4304
4305 assert!(
4306 bundle_on
4307 .capsules
4308 .iter()
4309 .any(|c| c.expansion_handle == "memory:m_b2"),
4310 "on-topic query (cosine=1.0 ≥ floor) must surface m_b2; \
4311 got capsules: {:?}",
4312 bundle_on
4313 .capsules
4314 .iter()
4315 .map(|c| &c.expansion_handle)
4316 .collect::<Vec<_>>()
4317 );
4318
4319 let conn3 = rusqlite::Connection::open_in_memory().expect("in-memory 3");
4322 crate::schema::initialize(&conn3).expect("init schema 3");
4323 insert_memory_with_embedding(
4324 &conn3,
4325 "m_b3",
4326 "cookie chocolate TOPIC_B recipe".to_string().as_str(),
4327 &emb,
4328 );
4329
4330 let bundle_noop_floor = retrieve_context_with_embedder(
4331 &conn3,
4332 "/fake-repo",
4333 &weights,
4334 ContextRequest {
4335 stage: "localization".to_string(),
4336 query: "cookie chocolate TOPIC_A".to_string(),
4338 budget_tokens: 4000,
4339 min_semantic_score: 0.0, ..Default::default()
4341 },
4342 &[],
4343 &emb,
4344 )
4345 .expect("retrieve noop floor");
4346
4347 assert!(
4349 bundle_noop_floor
4350 .capsules
4351 .iter()
4352 .any(|c| c.expansion_handle == "memory:m_b3"),
4353 "with floor=0.0 (disabled), off-topic-cosine memory must still surface via FTS; \
4354 got: {:?}",
4355 bundle_noop_floor
4356 .capsules
4357 .iter()
4358 .map(|c| &c.expansion_handle)
4359 .collect::<Vec<_>>()
4360 );
4361 }
4362
4363 #[cfg(feature = "embeddings")]
4392 #[test]
4393 fn d1f_token_economy_fewer_capsules_signal_preserved() {
4394 struct OracleTopicEmbedder;
4396 impl embeddings::Embedder for OracleTopicEmbedder {
4397 fn embed(&self, text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
4398 let mut v = vec![0.0f32; 8];
4399 if text.contains("TOPIC_A") {
4400 v[0] = 1.0; } else {
4402 v[1] = 1.0; }
4404 Ok(v)
4405 }
4406 fn model_id(&self) -> &str {
4407 "oracle-topic-d8"
4408 }
4409 fn dim(&self) -> usize {
4410 8
4411 }
4412 }
4413
4414 let oracle = OracleTopicEmbedder;
4415
4416 let setup = |conn: &rusqlite::Connection| {
4418 for (mid, text) in [
4421 ("m_dup1", "TOPIC_A prefer ripgrep for searching"),
4422 ("m_dup2", "TOPIC_A rg is the fastest searcher"),
4423 ("m_dup3", "TOPIC_A use rg tool to find patterns"),
4424 (
4426 "m_relevant",
4427 "TOPIC_A critical lesson about search performance",
4428 ),
4429 ("m_noise1", "chocolate cookie baking TOPIC_B recipe"),
4431 ("m_noise2", "gardening tulip planting TOPIC_B spring"),
4432 ] {
4433 insert_memory_with_embedding(conn, mid, text, &oracle);
4434 }
4435 };
4436
4437 let weights = kimetsu_core::config::BrokerWeights::default();
4438
4439 let conn_lean = rusqlite::Connection::open_in_memory().expect("in-memory lean");
4446 crate::schema::initialize(&conn_lean).expect("init schema lean");
4447 setup(&conn_lean);
4448
4449 let bundle_lean = retrieve_context_with_embedder(
4450 &conn_lean,
4451 "/fake-repo",
4452 &weights,
4453 ContextRequest {
4454 stage: "localization".to_string(),
4455 query: "TOPIC_A search performance".to_string(),
4456 budget_tokens: 20_000,
4457 min_semantic_score: 0.0, ..Default::default()
4459 },
4460 &[],
4461 &embeddings::NoopEmbedder,
4462 )
4463 .expect("retrieve lean");
4464
4465 let lean_count = bundle_lean
4466 .capsules
4467 .iter()
4468 .filter(|c| c.expansion_handle.starts_with("memory:"))
4469 .count();
4470
4471 let conn_emb = rusqlite::Connection::open_in_memory().expect("in-memory emb");
4473 crate::schema::initialize(&conn_emb).expect("init schema emb");
4474 setup(&conn_emb);
4475
4476 let bundle_emb = retrieve_context_with_embedder(
4477 &conn_emb,
4478 "/fake-repo",
4479 &weights,
4480 ContextRequest {
4481 stage: "localization".to_string(),
4482 query: "TOPIC_A search performance".to_string(),
4483 budget_tokens: 20_000,
4484 min_semantic_score: 0.5, ..Default::default()
4486 },
4487 &[],
4488 &oracle,
4489 )
4490 .expect("retrieve with embeddings");
4491
4492 let emb_count = bundle_emb
4493 .capsules
4494 .iter()
4495 .filter(|c| c.expansion_handle.starts_with("memory:"))
4496 .count();
4497
4498 assert!(
4500 emb_count < lean_count,
4501 "D1e must reduce capsule count: embedding path {emb_count} must be \
4502 < lean path {lean_count}. Embedding capsules: {:?}",
4503 bundle_emb
4504 .capsules
4505 .iter()
4506 .map(|c| &c.expansion_handle)
4507 .collect::<Vec<_>>()
4508 );
4509
4510 assert!(
4512 bundle_emb
4513 .capsules
4514 .iter()
4515 .any(|c| c.expansion_handle == "memory:m_relevant"),
4516 "m_relevant must survive D1e selection (signal preserved); \
4517 embedding capsules: {:?}",
4518 bundle_emb
4519 .capsules
4520 .iter()
4521 .map(|c| &c.expansion_handle)
4522 .collect::<Vec<_>>()
4523 );
4524
4525 let lean_tokens: u32 = bundle_lean.capsules.iter().map(|c| c.token_estimate).sum();
4527 let emb_tokens: u32 = bundle_emb.capsules.iter().map(|c| c.token_estimate).sum();
4528 assert!(
4529 emb_tokens < lean_tokens,
4530 "D1e must reduce token usage: emb={emb_tokens} must be < lean={lean_tokens}"
4531 );
4532 }
4533
4534 #[test]
4540 fn lean_noop_embedder_uses_fts_then_recency_unchanged() {
4541 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4544 crate::schema::initialize(&conn).expect("init schema");
4545
4546 for (mid, text) in [
4548 ("m_x", "use git rebase to clean history"),
4549 ("m_y", "grep finds text quickly"),
4550 ] {
4551 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4552 conn.execute(
4553 "INSERT INTO memories (
4554 memory_id, scope, kind, text, normalized_text, confidence,
4555 source_event_id, provenance_snapshot_json, created_at,
4556 use_count, usefulness_score
4557 )
4558 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
4559 '2026-01-01T00:00:00Z', 0, 0.0)",
4560 rusqlite::params![mid, text, normalized],
4561 )
4562 .expect("insert");
4563 conn.execute(
4564 "INSERT INTO memories_fts (memory_id, text, kind, scope) VALUES (?1, ?2, 'fact', 'global_user')",
4565 rusqlite::params![mid, text],
4566 )
4567 .expect("insert fts");
4568 }
4569
4570 let weights = kimetsu_core::config::BrokerWeights::default();
4571 let bundle = retrieve_context_with_embedder(
4573 &conn,
4574 "/fake-repo",
4575 &weights,
4576 ContextRequest {
4577 stage: "localization".to_string(),
4578 query: "grep text".to_string(),
4579 budget_tokens: 4000,
4580 ..Default::default()
4581 },
4582 &[],
4583 &embeddings::NoopEmbedder,
4584 )
4585 .expect("retrieve with NoopEmbedder must not panic");
4586
4587 let handles: Vec<&str> = bundle
4589 .capsules
4590 .iter()
4591 .filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
4592 .collect();
4593 assert!(
4594 handles.contains(&"m_y"),
4595 "m_y must surface via FTS on lean path; got {handles:?}"
4596 );
4597 }
4599
4600 #[test]
4606 fn classify_task_maps_each_kind_deterministically() {
4607 assert_eq!(
4609 classify_task("fix the panic in the parser"),
4610 TaskKind::Debug,
4611 "contains 'fix' and 'panic'"
4612 );
4613 assert_eq!(
4614 classify_task("there is a crash in auth when calling login"),
4615 TaskKind::Debug,
4616 "contains 'crash'"
4617 );
4618 assert_eq!(
4619 classify_task("debug the failing test"),
4620 TaskKind::Debug,
4621 "contains 'debug' and 'fail'"
4622 );
4623
4624 assert_eq!(
4626 classify_task("investigate why retrieval is slow"),
4627 TaskKind::Investigation,
4628 "contains 'investigate' and 'why'"
4629 );
4630 assert_eq!(
4631 classify_task("analyze the root cause of the latency"),
4632 TaskKind::Investigation,
4633 "contains 'analyze' and 'root cause'"
4634 );
4635
4636 assert_eq!(
4638 classify_task("refactor the auth module"),
4639 TaskKind::Refactor,
4640 "contains 'refactor'"
4641 );
4642 assert_eq!(
4643 classify_task("rename the config struct"),
4644 TaskKind::Refactor,
4645 "contains 'rename'"
4646 );
4647 assert_eq!(
4648 classify_task("simplify the retry handling logic"),
4649 TaskKind::Refactor,
4650 "contains 'simplify'"
4651 );
4652
4653 assert_eq!(
4655 classify_task("document the API endpoints"),
4656 TaskKind::Docs,
4657 "contains 'document'"
4658 );
4659 assert_eq!(
4660 classify_task("update the readme with new instructions"),
4661 TaskKind::Docs,
4662 "contains 'readme'"
4663 );
4664 assert_eq!(
4665 classify_task("add a docstring to the main function"),
4666 TaskKind::Docs,
4667 "contains 'docstring'"
4668 );
4669
4670 assert_eq!(
4672 classify_task("add a dark mode toggle"),
4673 TaskKind::Feature,
4674 "no debug/refactor/docs/investigate keyword"
4675 );
4676 assert_eq!(
4677 classify_task("implement the new caching layer"),
4678 TaskKind::Feature,
4679 "no debug/refactor/docs/investigate keyword"
4680 );
4681 assert_eq!(
4682 classify_task("build the export pipeline"),
4683 TaskKind::Feature,
4684 "no debug/refactor/docs/investigate keyword"
4685 );
4686 }
4687
4688 #[test]
4690 fn classify_task_respects_precedence_order() {
4691 assert_eq!(
4693 classify_task("fix and refactor the login module"),
4694 TaskKind::Debug,
4695 "Debug > Refactor"
4696 );
4697 assert_eq!(
4699 classify_task("investigate and refactor the cache layer"),
4700 TaskKind::Investigation,
4701 "Investigation > Refactor"
4702 );
4703 assert_eq!(
4705 classify_task("investigate the docs and document the API"),
4706 TaskKind::Investigation,
4707 "Investigation > Docs"
4708 );
4709 assert_eq!(
4711 classify_task("refactor and add docs"),
4712 TaskKind::Refactor,
4713 "Refactor > Docs"
4714 );
4715 assert_eq!(
4717 classify_task("fix the bug and investigate the regression"),
4718 TaskKind::Debug,
4719 "Debug > Investigation"
4720 );
4721 }
4722
4723 fn two_kinds_one_strong() -> Vec<Candidate> {
4728 let mk = |kind: &str, raw: f32| Candidate {
4729 capsule: ContextCapsule {
4730 id: format!("{kind}-1"),
4731 kind: kind.to_string(),
4732 summary: String::new(),
4733 token_estimate: 0,
4734 expansion_handle: String::new(),
4735 provenance: Vec::new(),
4736 confidence: 0.0,
4737 freshness: 0.0,
4738 relevance: 0.0,
4739 scope_weight: 0.0,
4740 score: 0.0,
4741 },
4742 raw_relevance: raw,
4743 embedding: None,
4744 cosine: None,
4745 created_at: None,
4746 };
4747 vec![mk("memory", 0.9), mk("repo_file", 0.1)]
4748 }
4749
4750 #[test]
4753 fn per_kind_normalization_flatters_the_best_of_a_weak_kind() {
4754 let mut candidates = two_kinds_one_strong();
4755 let weights = StageWeights {
4756 relevance: 1.0,
4757 confidence: 0.0,
4758 freshness: 0.0,
4759 scope: 0.0,
4760 };
4761 normalize_and_score(&mut candidates, weights, Normalization::PerKind);
4762 assert!((candidates[0].capsule.relevance - 1.0).abs() < 1e-6);
4763 assert!(
4764 (candidates[1].capsule.relevance - 1.0).abs() < 1e-6,
4765 "per-kind gives the lone weak repo_file relevance 1.0, got {}",
4766 candidates[1].capsule.relevance
4767 );
4768 }
4769
4770 #[test]
4773 fn global_normalization_keeps_relevance_comparable_across_kinds() {
4774 let mut candidates = two_kinds_one_strong();
4775 let weights = StageWeights {
4776 relevance: 1.0,
4777 confidence: 0.0,
4778 freshness: 0.0,
4779 scope: 0.0,
4780 };
4781 normalize_and_score(&mut candidates, weights, Normalization::Global);
4782 assert!((candidates[0].capsule.relevance - 1.0).abs() < 1e-6);
4783 let weak = candidates[1].capsule.relevance;
4784 assert!(
4785 (weak - (0.1 / 0.9)).abs() < 1e-6,
4786 "global normalizes against the single max, got {weak}"
4787 );
4788 assert!(weak < candidates[0].capsule.relevance);
4789 }
4790
4791 #[test]
4794 fn unknown_normalization_falls_back_to_per_kind() {
4795 assert_eq!(Normalization::from_config(""), Normalization::PerKind);
4796 assert_eq!(
4797 Normalization::from_config("per_kind"),
4798 Normalization::PerKind
4799 );
4800 assert_eq!(
4801 Normalization::from_config("nonsense"),
4802 Normalization::PerKind
4803 );
4804 assert_eq!(Normalization::from_config("global"), Normalization::Global);
4805 assert_eq!(
4806 Normalization::from_config(" GLOBAL "),
4807 Normalization::Global
4808 );
4809 }
4810
4811 #[test]
4812 fn weights_for_task_kind_renormalizes_to_unit_sum() {
4813 let base = StageWeights {
4814 relevance: 0.50,
4815 confidence: 0.20,
4816 freshness: 0.20,
4817 scope: 0.10,
4818 };
4819 let original_sum = base.relevance + base.confidence + base.freshness + base.scope;
4820
4821 for kind in [
4822 TaskKind::Debug,
4823 TaskKind::Refactor,
4824 TaskKind::Investigation,
4825 TaskKind::Docs,
4826 ] {
4827 let w = weights_for_task_kind(base.clone(), kind);
4828 let new_sum = w.relevance + w.confidence + w.freshness + w.scope;
4829 assert!(
4831 (new_sum - original_sum).abs() < 1e-4,
4832 "weights_for_task_kind({kind:?}) sum {new_sum} differs from {original_sum}"
4833 );
4834 }
4835 }
4836
4837 #[test]
4839 fn weights_for_task_kind_feature_is_unchanged() {
4840 let base = StageWeights {
4841 relevance: 0.40,
4842 confidence: 0.30,
4843 freshness: 0.20,
4844 scope: 0.10,
4845 };
4846 let w = weights_for_task_kind(base.clone(), TaskKind::Feature);
4847 assert!((w.relevance - base.relevance).abs() < f32::EPSILON);
4848 assert!((w.confidence - base.confidence).abs() < f32::EPSILON);
4849 assert!((w.freshness - base.freshness).abs() < f32::EPSILON);
4850 assert!((w.scope - base.scope).abs() < f32::EPSILON);
4851 }
4852
4853 #[test]
4856 fn weights_for_task_kind_debug_up_freshness_fraction() {
4857 let base = StageWeights {
4858 relevance: 0.50,
4859 confidence: 0.20,
4860 freshness: 0.20,
4861 scope: 0.10,
4862 };
4863 let debug_w = weights_for_task_kind(base.clone(), TaskKind::Debug);
4864 assert!(
4866 debug_w.freshness > base.freshness,
4867 "Debug must increase freshness fraction: {debug_w:?}"
4868 );
4869 }
4870
4871 #[test]
4874 fn weights_for_task_kind_refactor_up_scope_fraction() {
4875 let base = StageWeights {
4876 relevance: 0.50,
4877 confidence: 0.20,
4878 freshness: 0.20,
4879 scope: 0.10,
4880 };
4881 let refactor_w = weights_for_task_kind(base.clone(), TaskKind::Refactor);
4882 assert!(
4883 refactor_w.scope > base.scope,
4884 "Refactor must increase scope fraction: {refactor_w:?}"
4885 );
4886 }
4887
4888 #[test]
4891 fn task_kind_feature_is_retrieval_neutral() {
4892 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4893 crate::schema::initialize(&conn).expect("init schema");
4894
4895 for (mid, db_kind, text) in [
4899 ("m1", "failure_pattern", "linker not found error in build"),
4900 ("m2", "convention", "use snake_case for all identifiers"),
4901 ("m3", "fact", "the cache is invalidated on every deploy"),
4902 ] {
4903 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4904 conn.execute(
4905 "INSERT INTO memories (
4906 memory_id, scope, kind, text, normalized_text, confidence,
4907 source_event_id, provenance_snapshot_json, created_at,
4908 use_count, usefulness_score
4909 )
4910 VALUES (?1, 'project', ?2, ?3, ?4, 1.0, NULL, '{}',
4911 '2026-01-01T00:00:00Z', 0, 0.0)",
4912 rusqlite::params![mid, db_kind, text, normalized],
4913 )
4914 .expect("insert memory");
4915 conn.execute(
4916 "INSERT INTO memories_fts (memory_id, text, kind, scope)
4917 VALUES (?1, ?2, ?3, 'project')",
4918 rusqlite::params![mid, text, db_kind],
4919 )
4920 .expect("insert fts");
4921 }
4922
4923 let weights = kimetsu_core::config::BrokerWeights::default();
4924 let query = "cache convention failure".to_string();
4925
4926 let baseline = retrieve_context_with_embedder(
4928 &conn,
4929 "/fake-repo",
4930 &weights,
4931 ContextRequest {
4932 stage: "localization".to_string(),
4933 query: query.clone(),
4934 budget_tokens: 4000,
4935 ..Default::default()
4936 },
4937 &[],
4938 &embeddings::NoopEmbedder,
4939 )
4940 .expect("baseline retrieve");
4941
4942 let feature = retrieve_context_with_embedder(
4944 &conn,
4945 "/fake-repo",
4946 &weights,
4947 ContextRequest {
4948 stage: "localization".to_string(),
4949 query: query.clone(),
4950 budget_tokens: 4000,
4951 task_kind: TaskKind::Feature,
4952 ..Default::default()
4953 },
4954 &[],
4955 &embeddings::NoopEmbedder,
4956 )
4957 .expect("feature retrieve");
4958
4959 let baseline_ids: Vec<&str> = baseline
4960 .capsules
4961 .iter()
4962 .map(|c| c.expansion_handle.as_str())
4963 .collect();
4964 let feature_ids: Vec<&str> = feature
4965 .capsules
4966 .iter()
4967 .map(|c| c.expansion_handle.as_str())
4968 .collect();
4969 assert_eq!(
4970 baseline_ids, feature_ids,
4971 "task_kind=Feature must produce identical retrieval to default; \
4972 baseline={baseline_ids:?} feature={feature_ids:?}"
4973 );
4974
4975 let baseline_scores: Vec<f32> = baseline.capsules.iter().map(|c| c.score).collect();
4976 let feature_scores: Vec<f32> = feature.capsules.iter().map(|c| c.score).collect();
4977 for (b, f) in baseline_scores.iter().zip(feature_scores.iter()) {
4978 assert!(
4979 (b - f).abs() < 1e-5,
4980 "scores must be identical: baseline={b} feature={f}"
4981 );
4982 }
4983 }
4984
4985 #[test]
4996 fn debug_surfaces_more_failure_pattern_than_docs() {
4997 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4998 crate::schema::initialize(&conn).expect("init schema");
4999
5000 for (i, text) in [
5004 "auth token expired causes login failure",
5005 "auth service crash on null pointer",
5006 "auth regression after upgrade breaks sessions",
5007 "auth error when certificate is invalid",
5008 ]
5009 .iter()
5010 .enumerate()
5011 {
5012 let mid = format!("mfp{i}");
5013 let normalized = kimetsu_core::memory::normalize_memory_text(text);
5014 conn.execute(
5015 "INSERT INTO memories (
5016 memory_id, scope, kind, text, normalized_text, confidence,
5017 source_event_id, provenance_snapshot_json, created_at,
5018 use_count, usefulness_score
5019 )
5020 VALUES (?1, 'project', 'failure_pattern', ?2, ?3, 1.0, NULL, '{}',
5021 '2026-01-01T00:00:00Z', 0, 0.0)",
5022 rusqlite::params![mid, text, normalized],
5023 )
5024 .expect("insert failure_pattern");
5025 conn.execute(
5026 "INSERT INTO memories_fts (memory_id, text, kind, scope)
5027 VALUES (?1, ?2, 'failure_pattern', 'project')",
5028 rusqlite::params![mid, text],
5029 )
5030 .expect("insert fts");
5031 }
5032
5033 for (i, (db_kind, text)) in [
5036 ("convention", "auth module uses bearer tokens by convention"),
5037 ("convention", "auth scopes are documented in the API guide"),
5038 ("fact", "auth service runs on port 8443 in production"),
5039 ("fact", "auth uses JWT with RS256 signing for all tokens"),
5040 ]
5041 .iter()
5042 .enumerate()
5043 {
5044 let mid = format!("mconv{i}");
5045 let normalized = kimetsu_core::memory::normalize_memory_text(text);
5046 conn.execute(
5047 "INSERT INTO memories (
5048 memory_id, scope, kind, text, normalized_text, confidence,
5049 source_event_id, provenance_snapshot_json, created_at,
5050 use_count, usefulness_score
5051 )
5052 VALUES (?1, 'project', ?2, ?3, ?4, 1.0, NULL, '{}',
5053 '2026-01-01T00:00:00Z', 0, 0.0)",
5054 rusqlite::params![mid, db_kind, text, normalized],
5055 )
5056 .expect("insert convention/fact");
5057 conn.execute(
5058 "INSERT INTO memories_fts (memory_id, text, kind, scope)
5059 VALUES (?1, ?2, ?3, 'project')",
5060 rusqlite::params![mid, text, db_kind],
5061 )
5062 .expect("insert fts");
5063 }
5064
5065 let weights = kimetsu_core::config::BrokerWeights::default();
5066 let query = "auth token failure".to_string();
5067
5068 let debug_bundle = retrieve_context_with_embedder(
5070 &conn,
5071 "/fake-repo",
5072 &weights,
5073 ContextRequest {
5074 stage: "localization".to_string(),
5075 query: query.clone(),
5076 budget_tokens: 4000,
5077 max_capsules: 4,
5078 task_kind: TaskKind::Debug,
5079 ..Default::default()
5080 },
5081 &[],
5082 &embeddings::NoopEmbedder,
5083 )
5084 .expect("debug retrieve");
5085
5086 let docs_bundle = retrieve_context_with_embedder(
5088 &conn,
5089 "/fake-repo",
5090 &weights,
5091 ContextRequest {
5092 stage: "localization".to_string(),
5093 query: query.clone(),
5094 budget_tokens: 4000,
5095 max_capsules: 4,
5096 task_kind: TaskKind::Docs,
5097 ..Default::default()
5098 },
5099 &[],
5100 &embeddings::NoopEmbedder,
5101 )
5102 .expect("docs retrieve");
5103
5104 let count_failure_pattern = |bundle: &ContextBundle| -> usize {
5107 bundle
5108 .capsules
5109 .iter()
5110 .filter(|c| capsule_matches_kind(c, "failure_pattern"))
5111 .count()
5112 };
5113
5114 let debug_fp = count_failure_pattern(&debug_bundle);
5115 let docs_fp = count_failure_pattern(&docs_bundle);
5116
5117 assert!(
5118 debug_fp > docs_fp,
5119 "Debug must surface strictly more failure_pattern capsules than Docs: \
5120 debug_fp={debug_fp} docs_fp={docs_fp}\n\
5121 Debug capsules: {:?}\n\
5122 Docs capsules: {:?}",
5123 debug_bundle
5124 .capsules
5125 .iter()
5126 .map(|c| format!("{}:{}", c.kind, &c.summary[..c.summary.len().min(60)]))
5127 .collect::<Vec<_>>(),
5128 docs_bundle
5129 .capsules
5130 .iter()
5131 .map(|c| format!("{}:{}", c.kind, &c.summary[..c.summary.len().min(60)]))
5132 .collect::<Vec<_>>(),
5133 );
5134 }
5135
5136 fn init_db_with_memory(memory_id: &str, text: &str) -> rusqlite::Connection {
5139 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
5140 crate::schema::initialize(&conn).expect("init schema");
5141 let normalized = kimetsu_core::memory::normalize_memory_text(text);
5142 conn.execute(
5143 "INSERT INTO memories (
5144 memory_id, scope, kind, text, normalized_text, confidence,
5145 source_event_id, provenance_snapshot_json, created_at,
5146 use_count, usefulness_score
5147 )
5148 VALUES (?1, 'project', 'fact', ?2, ?3, 1.0, NULL, '{}',
5149 '2026-01-01T00:00:00Z', 0, 0.0)",
5150 rusqlite::params![memory_id, text, normalized],
5151 )
5152 .expect("insert memory");
5153 conn
5154 }
5155
5156 #[test]
5158 fn resolve_capsule_memory_returns_full_text() {
5159 let conn = init_db_with_memory("test-mem-id", "Use rg over grep for speed");
5160 let repo_root = std::path::Path::new("/fake-repo");
5161 let result =
5162 resolve_capsule(&conn, repo_root, "memory:test-mem-id").expect("should resolve");
5163 assert_eq!(result, "Use rg over grep for speed");
5164 }
5165
5166 #[test]
5168 fn resolve_capsule_memory_missing_id_returns_err() {
5169 let conn = init_db_with_memory("real-id", "some text");
5170 let repo_root = std::path::Path::new("/fake-repo");
5171 let err = resolve_capsule(&conn, repo_root, "memory:nonexistent-id")
5172 .expect_err("should error for missing memory");
5173 assert!(
5174 err.to_string().contains("no active memory"),
5175 "error message should mention missing: {err}"
5176 );
5177 }
5178
5179 #[test]
5181 fn resolve_capsule_file_returns_bounded_content() {
5182 let dir = make_test_dir("f2_file_resolve");
5183 let content = "hello from the file\n";
5184 std::fs::write(dir.join("notes.txt"), content).expect("write");
5185 let result = resolve_capsule(
5186 &rusqlite::Connection::open_in_memory().expect("open"),
5188 &dir,
5189 "file:notes.txt",
5190 )
5191 .expect("should resolve file");
5192 assert!(result.contains("hello from the file"));
5193 std::fs::remove_dir_all(&dir).ok();
5194 }
5195
5196 #[test]
5198 fn resolve_capsule_file_caps_large_file() {
5199 let dir = make_test_dir("f2_file_cap");
5200 let big = "A".repeat(FILE_EXPAND_CAP_BYTES * 3);
5201 std::fs::write(dir.join("big.txt"), &big).expect("write");
5202 let result = resolve_capsule(
5203 &rusqlite::Connection::open_in_memory().expect("open"),
5204 &dir,
5205 "file:big.txt",
5206 )
5207 .expect("should resolve large file");
5208 assert!(
5209 result.len() <= FILE_EXPAND_CAP_BYTES + 200,
5210 "result should be bounded: got {} bytes",
5211 result.len()
5212 );
5213 assert!(
5214 result.contains("truncated"),
5215 "truncation marker should be present"
5216 );
5217 std::fs::remove_dir_all(&dir).ok();
5218 }
5219
5220 #[test]
5222 fn resolve_capsule_unknown_handle_returns_err() {
5223 let conn = rusqlite::Connection::open_in_memory().expect("open");
5224 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "blob:abc123")
5225 .expect_err("should error");
5226 assert!(
5227 err.to_string().contains("unrecognised handle"),
5228 "got: {err}"
5229 );
5230 }
5231
5232 #[test]
5234 fn resolve_capsule_malformed_handle_returns_err() {
5235 let conn = rusqlite::Connection::open_in_memory().expect("open");
5236 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "justnocolon")
5237 .expect_err("should error");
5238 assert!(
5239 err.to_string().contains("unrecognised handle"),
5240 "got: {err}"
5241 );
5242 }
5243
5244 #[test]
5246 fn resolve_capsule_run_handle_returns_deferred_err() {
5247 let conn = rusqlite::Connection::open_in_memory().expect("open");
5248 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "run:some-run-id")
5249 .expect_err("run: should be deferred err");
5250 assert!(err.to_string().contains("not yet supported"), "got: {err}");
5251 }
5252
5253 #[test]
5255 fn resolve_capsule_file_rejects_absolute_path() {
5256 let conn = rusqlite::Connection::open_in_memory().expect("open");
5257 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "file:/etc/passwd")
5258 .expect_err("should reject absolute path");
5259 assert!(err.to_string().contains("absolute path"), "got: {err}");
5260 }
5261
5262 fn make_capsule(summary: &str, score: f32) -> ContextCapsule {
5265 ContextCapsule {
5266 id: new_id().to_string(),
5267 kind: "memory".to_string(),
5268 summary: summary.to_string(),
5269 token_estimate: 10,
5270 expansion_handle: format!("memory:{}", new_id()),
5271 provenance: vec![],
5272 confidence: 1.0,
5273 freshness: 1.0,
5274 relevance: 1.0,
5275 scope_weight: 1.0,
5276 score,
5277 }
5278 }
5279
5280 #[test]
5283 fn rerank_capsules_reorders_by_query_overlap() {
5284 use crate::embeddings::StubReranker;
5285
5286 let query = "rust async tokio";
5289 let high_overlap = make_capsule("rust async tokio runtime", 0.0);
5290 let low_overlap = make_capsule("python django framework", 0.0);
5291 let capsules = vec![low_overlap.clone(), high_overlap.clone()];
5293
5294 let ranked = rerank_capsules(query, capsules, &StubReranker, 0.0, 0);
5295
5296 assert_eq!(ranked.len(), 2, "both capsules should survive (floor=0)");
5297 assert!(
5299 ranked[0].summary.contains("rust"),
5300 "rust capsule must be first, got: {:?}",
5301 ranked[0].summary
5302 );
5303 assert!(
5305 ranked[0].score > 0.05,
5306 "score must be overwritten by reranker: {}",
5307 ranked[0].score
5308 );
5309 assert!(
5311 ranked[0].score > ranked[1].score,
5312 "high overlap must score higher: {} vs {}",
5313 ranked[0].score,
5314 ranked[1].score
5315 );
5316 }
5317
5318 #[test]
5322 fn rerank_capsules_floor_drops_zero_overlap() {
5323 use crate::embeddings::StubReranker;
5324
5325 let query = "rust async tokio";
5326 let high = make_capsule("rust async tokio runtime", 0.0);
5327 let zero = make_capsule("completely unrelated document xyz", 0.0); let capsules = vec![high, zero];
5330 let ranked = rerank_capsules(query, capsules, &StubReranker, 0.3, 0);
5331
5332 assert_eq!(ranked.len(), 1, "zero-overlap capsule must be dropped");
5334 assert!(
5335 ranked[0].summary.contains("rust"),
5336 "only rust capsule should survive"
5337 );
5338 }
5339
5340 #[test]
5342 fn rerank_capsules_cap_truncates() {
5343 use crate::embeddings::StubReranker;
5344
5345 let query = "alpha beta gamma";
5346 let capsules = vec![
5347 make_capsule("alpha beta gamma delta", 0.0),
5348 make_capsule("alpha beta", 0.0),
5349 make_capsule("alpha", 0.0),
5350 make_capsule("unrelated xyz", 0.0),
5351 ];
5352
5353 let ranked = rerank_capsules(query, capsules, &StubReranker, 0.0, 2);
5354 assert_eq!(ranked.len(), 2, "cap=2 must truncate to 2 results");
5355 assert!(
5357 ranked[0].score >= ranked[1].score,
5358 "results must be sorted descending"
5359 );
5360 }
5361
5362 #[test]
5364 fn rerank_capsules_fail_open_preserves_input_order() {
5365 struct FailingReranker;
5366 impl crate::embeddings::Reranker for FailingReranker {
5367 fn rerank(
5368 &self,
5369 _query: &str,
5370 _docs: &[&str],
5371 ) -> Result<Vec<f32>, crate::embeddings::EmbedderError> {
5372 Err(crate::embeddings::EmbedderError::EmbedFailed(
5373 "simulated failure".into(),
5374 ))
5375 }
5376 fn model_id(&self) -> &str {
5377 "fail-reranker"
5378 }
5379 }
5380
5381 let query = "anything";
5382 let c1 = make_capsule("first capsule", 0.9);
5383 let c2 = make_capsule("second capsule", 0.5);
5384 let c3 = make_capsule("third capsule", 0.1);
5385 let capsules = vec![c1.clone(), c2.clone(), c3.clone()];
5386
5387 let out = rerank_capsules(query, capsules, &FailingReranker, 0.0, 0);
5388
5389 assert_eq!(out.len(), 3, "all capsules must be returned on error");
5391 assert_eq!(out[0].summary, c1.summary, "order must be preserved");
5392 assert_eq!(out[1].summary, c2.summary, "order must be preserved");
5393 assert_eq!(out[2].summary, c3.summary, "order must be preserved");
5394 }
5395
5396 #[test]
5398 fn rerank_capsules_empty_input_returns_empty() {
5399 use crate::embeddings::StubReranker;
5400 let out = rerank_capsules("query", vec![], &StubReranker, 0.0, 0);
5401 assert!(out.is_empty());
5402 }
5403
5404 #[test]
5408 fn compress_for_render_short_text_unchanged() {
5409 let text = "project:fact - Use cargo fmt before committing.";
5410 let out = compress_for_render(text, 3);
5411 assert_eq!(out, text, "short text must not be altered");
5412 }
5413
5414 #[test]
5416 fn compress_for_render_strips_tags_prefix() {
5417 let text = "[tags: rust, cargo] Always run cargo clippy before submitting a PR.";
5418 let out = compress_for_render(text, 3);
5419 assert!(
5420 !out.starts_with('['),
5421 "tags prefix must be stripped, got: {out:?}"
5422 );
5423 assert!(
5424 out.contains("cargo clippy"),
5425 "body must remain, got: {out:?}"
5426 );
5427 }
5428
5429 #[test]
5431 fn compress_for_render_strips_context_suffix() {
5432 let text =
5433 "project:fact - Use cargo fmt. Always clippy clean. (context: Kimetsu brain lesson)";
5434 let out = compress_for_render(text, 5);
5435 assert!(
5436 !out.contains("(context:"),
5437 "context suffix must be stripped, got: {out:?}"
5438 );
5439 assert!(out.contains("cargo fmt"), "body must remain, got: {out:?}");
5440 }
5441
5442 #[test]
5444 fn compress_for_render_caps_sentences() {
5445 let text =
5446 "project:fact - First sentence. Second sentence. Third sentence. Fourth sentence.";
5447 let out = compress_for_render(text, 2);
5448 assert!(out.contains("First"), "first sentence must be present");
5450 assert!(out.contains("Second"), "second sentence must be present");
5451 assert!(
5452 !out.contains("Third"),
5453 "third sentence must be truncated, got: {out:?}"
5454 );
5455 }
5456
5457 #[test]
5459 fn compress_for_render_preserves_scope_prefix() {
5460 let text = "global_user:convention - First rule. Second rule. Third rule. Fourth rule.";
5461 let out = compress_for_render(text, 2);
5462 assert!(
5463 out.starts_with("global_user:convention - "),
5464 "scope prefix must be preserved, got: {out:?}"
5465 );
5466 assert!(out.contains("First"), "first sentence must remain");
5467 assert!(!out.contains("Third"), "third sentence must be truncated");
5468 }
5469
5470 #[test]
5472 fn compress_for_render_empty_input_safe() {
5473 let out = compress_for_render("", 3);
5474 assert_eq!(out, "", "empty input must return empty string");
5475 }
5476
5477 #[test]
5479 fn compress_for_render_zero_max_sentences_returns_original() {
5480 let text = "project:fact - Some lesson that is quite long. It keeps going. And going.";
5481 let out = compress_for_render(text, 0);
5482 assert_eq!(out, text);
5483 }
5484
5485 #[test]
5487 fn compress_for_render_utf8_safe() {
5488 let text = "project:fact - こんにちは世界. Hello world. Third sentence. Fourth sentence.";
5489 let out = compress_for_render(text, 2);
5491 assert!(!out.is_empty(), "UTF-8 text must not produce empty output");
5492 assert!(!out.contains("Third"), "third sentence must be truncated");
5494 }
5495
5496 #[test]
5499 fn compress_for_render_long_memory_reduces_tokens_by_25_percent() {
5500 let long_summary = "project:fact - When a SQLite WAL file exists from a crashed process, \
5502 opening the DB causes the WAL to be replayed. The replayed WAL may contain \
5503 partial writes that corrupt the DB. Always check for WAL files before opening. \
5504 Delete the WAL only after verifying the DB is consistent. Use PRAGMA integrity_check \
5505 to validate after opening. If integrity_check fails, restore from backup. Never \
5506 truncate the WAL without replaying it first. This pattern applies to any \
5507 crash-recovery scenario.";
5508
5509 let raw_tokens = estimate_tokens(long_summary);
5510 assert!(
5511 raw_tokens > 60,
5512 "test precondition: raw memory must be >60 tokens, got {raw_tokens}"
5513 );
5514
5515 let compressed = compress_for_render(long_summary, 3);
5516 let compressed_tokens = estimate_tokens(&compressed);
5517
5518 let reduction = 1.0 - (compressed_tokens as f64 / raw_tokens as f64);
5519 assert!(
5520 reduction >= 0.25,
5521 "compression must reduce tokens by >=25% on long memories; \
5522 raw={raw_tokens} compressed={compressed_tokens} reduction={reduction:.2}"
5523 );
5524 }
5525}
5526
5527#[cfg(test)]
5528mod evidence_tests {
5529 use super::*;
5530
5531 fn conn_with(texts: &[&str]) -> Connection {
5532 let conn = Connection::open_in_memory().expect("open");
5533 crate::schema::initialize(&conn).expect("schema");
5534 for (i, text) in texts.iter().enumerate() {
5535 conn.execute(
5536 "INSERT INTO memories
5537 (memory_id, scope, kind, text, normalized_text, confidence,
5538 provenance_snapshot_json, created_at)
5539 VALUES (?1, 'project', 'fact', ?2, ?2, 0.9, '{}', '2026-01-01T00:00:00Z')",
5540 rusqlite::params![format!("m{i}"), text],
5541 )
5542 .expect("insert");
5543 }
5544 conn
5545 }
5546
5547 fn capsule(summary: &str) -> ContextCapsule {
5548 ContextCapsule {
5549 id: String::new(),
5550 kind: "memory".to_string(),
5551 summary: summary.to_string(),
5552 token_estimate: 10,
5553 expansion_handle: format!("memory:{summary}"),
5554 provenance: Vec::new(),
5555 confidence: 0.9,
5556 freshness: 0.5,
5557 relevance: 0.0,
5558 scope_weight: 0.9,
5559 score: 0.5,
5560 }
5561 }
5562
5563 fn bundle(capsules: Vec<ContextCapsule>, coverage: f32, uncovered: &[&str]) -> ContextBundle {
5564 ContextBundle {
5565 stage: "localization".to_string(),
5566 budget_tokens: 2000,
5567 used_tokens: 20,
5568 capsules,
5569 excluded: Vec::new(),
5570 skipped: false,
5571 top_score: 0.7,
5572 evidence_coverage: coverage,
5573 uncovered_terms: uncovered.iter().map(|s| s.to_string()).collect(),
5574 chronological: false,
5575 }
5576 }
5577
5578 #[test]
5581 fn full_coverage_names_nothing() {
5582 let conn = conn_with(&[
5583 "checkpoint the wal before copying brain.db",
5584 "vacuum reclaims dead pages",
5585 ]);
5586 let (coverage, uncovered) = evidence_coverage(
5587 &conn,
5588 "checkpoint wal",
5589 &[capsule(
5590 "project:fact - checkpoint the wal before copying brain.db",
5591 )],
5592 );
5593 assert!(coverage > 0.99, "got {coverage}");
5594 assert!(uncovered.is_empty(), "got {uncovered:?}");
5595 }
5596
5597 #[test]
5601 fn partial_coverage_names_the_missing_terms() {
5602 let conn = conn_with(&[
5603 "checkpoint the wal before copying brain.db",
5604 "the migration runner snapshots before each step",
5605 ]);
5606 let (coverage, uncovered) = evidence_coverage(
5607 &conn,
5608 "checkpoint wal migration",
5609 &[capsule(
5610 "project:fact - checkpoint the wal before copying brain.db",
5611 )],
5612 );
5613 assert!(coverage < 1.0, "coverage should be partial: {coverage}");
5614 assert!(
5615 uncovered.iter().any(|t| t.starts_with("migrat")),
5616 "the uncovered term must be named: {uncovered:?}"
5617 );
5618 }
5619
5620 #[test]
5623 fn coverage_is_collective_not_per_capsule() {
5624 let conn = conn_with(&[
5625 "checkpoint the wal before copying brain.db",
5626 "the migration runner snapshots before each step",
5627 ]);
5628 let (coverage, uncovered) = evidence_coverage(
5629 &conn,
5630 "checkpoint migration",
5631 &[
5632 capsule("project:fact - checkpoint the wal before copying"),
5633 capsule("project:fact - the migration runner snapshots first"),
5634 ],
5635 );
5636 assert!(
5637 coverage > 0.99,
5638 "neither capsule covers both terms, but together they do: {coverage}"
5639 );
5640 assert!(uncovered.is_empty(), "got {uncovered:?}");
5641 }
5642
5643 #[test]
5646 fn an_unmeasurable_query_does_not_claim_a_gap() {
5647 let conn = conn_with(&["checkpoint the wal"]);
5648 let (coverage, uncovered) =
5649 evidence_coverage(&conn, "the and of", &[capsule("project:fact - checkpoint")]);
5650 assert_eq!(coverage, 1.0);
5651 assert!(uncovered.is_empty());
5652 }
5653
5654 #[test]
5661 fn a_term_the_corpus_has_never_seen_counts_as_a_gap() {
5662 let conn = conn_with(&[
5663 "checkpoint the wal before copying brain.db",
5664 "vacuum reclaims dead pages",
5665 ]);
5666 let (coverage, uncovered) = evidence_coverage(
5667 &conn,
5668 "checkpoint the wal during a kubernetes rollout",
5669 &[capsule(
5670 "project:fact - checkpoint the wal before copying brain.db",
5671 )],
5672 );
5673 assert!(
5674 coverage <= PARTIAL_EVIDENCE_COVERAGE,
5675 "an unknown half of the question must read as thin, not complete: {coverage}"
5676 );
5677 assert!(
5678 uncovered.iter().any(|t| t.starts_with("kubernet")),
5679 "the unknown term must be named: {uncovered:?}"
5680 );
5681 }
5682
5683 #[test]
5686 fn a_ubiquitous_term_carries_no_weight() {
5687 let conn = conn_with(&["kimetsu checkpoint wal", "kimetsu vacuum pages"]);
5688 let (coverage, _) = evidence_coverage(
5689 &conn,
5690 "kimetsu vacuum",
5691 &[capsule("project:fact - kimetsu vacuum pages")],
5692 );
5693 assert!(coverage > 0.99, "got {coverage}");
5694 }
5695
5696 #[test]
5697 fn an_empty_query_does_not_claim_a_gap() {
5698 let conn = conn_with(&["checkpoint the wal"]);
5699 assert_eq!(evidence_coverage(&conn, "", &[]).0, 1.0);
5700 }
5701
5702 #[test]
5705 fn a_complete_bundle_gets_no_notice() {
5706 assert!(partial_evidence_notice(&bundle(vec![capsule("a")], 1.0, &[])).is_none());
5707 assert!(
5708 partial_evidence_notice(&bundle(vec![capsule("a")], 0.9, &["x"])).is_none(),
5709 "above the threshold is not partial"
5710 );
5711 }
5712
5713 #[test]
5714 fn an_empty_or_skipped_bundle_gets_no_notice() {
5715 let mut skipped = bundle(Vec::new(), 0.0, &["x"]);
5716 skipped.skipped = true;
5717 assert!(
5718 partial_evidence_notice(&skipped).is_none(),
5719 "an empty bundle already says everything it can"
5720 );
5721 assert!(partial_evidence_notice(&bundle(Vec::new(), 0.0, &["x"])).is_none());
5722 }
5723
5724 #[test]
5725 fn a_partial_bundle_names_what_is_missing_and_tells_the_reader_what_to_do() {
5726 let notice =
5727 partial_evidence_notice(&bundle(vec![capsule("a")], 0.3, &["migration", "rollback"]))
5728 .expect("a thin bundle must be flagged");
5729 assert!(notice.contains("migration"), "got: {notice}");
5730 assert!(notice.contains("rollback"), "got: {notice}");
5731 assert!(
5732 notice.contains("unknown"),
5733 "the notice must tell the reader to abstain, not just report a gap: {notice}"
5734 );
5735 }
5736
5737 #[test]
5739 fn the_notice_caps_how_many_terms_it_names() {
5740 let terms: Vec<String> = (0..12).map(|i| format!("term{i}")).collect();
5741 let refs: Vec<&str> = terms.iter().map(String::as_str).collect();
5742 let notice =
5743 partial_evidence_notice(&bundle(vec![capsule("a")], 0.1, &refs)).expect("flagged");
5744 assert!(notice.contains("and 6 more"), "got: {notice}");
5745 assert!(!notice.contains("term9"), "got: {notice}");
5746 }
5747
5748 #[test]
5754 fn the_y_ies_pair_shares_a_stem() {
5755 for (a, b) in [
5756 ("retry", "retries"),
5757 ("query", "queries"),
5758 ("policy", "policies"),
5759 ("memory", "memories"),
5760 ("binary", "binaries"),
5761 ("registry", "registries"),
5762 ] {
5763 assert_eq!(
5764 light_stem(a),
5765 light_stem(b),
5766 "{a}/{b} stemmed to {:?}/{:?}",
5767 light_stem(a),
5768 light_stem(b)
5769 );
5770 }
5771 }
5772
5773 #[test]
5775 fn a_vowel_y_is_not_stripped() {
5776 assert_eq!(light_stem("delay"), "delay");
5777 assert_eq!(light_stem("gateway"), "gateway");
5778 assert_eq!(light_stem("journeys"), "journey");
5781 }
5782
5783 #[test]
5786 fn short_words_keep_their_ending() {
5787 assert_eq!(light_stem("body"), "body");
5788 assert_eq!(light_stem("copy"), "copy");
5789 }
5790
5791 #[test]
5793 fn the_original_suffix_rules_still_hold() {
5794 assert_eq!(light_stem("benchmarked"), "benchmark");
5795 assert_eq!(light_stem("benchmarking"), "benchmark");
5796 assert_eq!(light_stem("migrations"), "migration");
5797 assert_eq!(light_stem("run"), "run");
5798 }
5799
5800 #[test]
5803 fn an_inflected_corpus_term_counts_as_covered() {
5804 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
5805 crate::schema::initialize(&conn).expect("init schema");
5806 let text = "the ingest worker retries a failed batch three times before giving up";
5807 let normalized = kimetsu_core::memory::normalize_memory_text(text);
5808 conn.execute(
5809 "
5810 INSERT INTO memories (
5811 memory_id, scope, kind, text, normalized_text, confidence,
5812 source_event_id, provenance_snapshot_json, created_at
5813 )
5814 VALUES ('m_retry', 'project', 'fact', ?1, ?2, 1.0, NULL, '{}',
5815 '2026-01-01T00:00:00Z')
5816 ",
5817 rusqlite::params![text, normalized],
5818 )
5819 .expect("insert memory");
5820 conn.execute(
5821 "INSERT INTO memories_fts (memory_id, text, kind, scope)
5822 VALUES ('m_retry', ?1, 'fact', 'project')",
5823 rusqlite::params![text],
5824 )
5825 .expect("insert fts");
5826
5827 let bundle = retrieve_context_with_embedder(
5828 &conn,
5829 "/fake-repo",
5830 &kimetsu_core::config::BrokerWeights::default(),
5831 ContextRequest {
5832 stage: "localization".to_string(),
5833 query: "how many times does the ingest worker retry a failed batch".to_string(),
5834 budget_tokens: 4000,
5835 ..Default::default()
5836 },
5837 &[],
5838 &embeddings::NoopEmbedder,
5839 )
5840 .expect("retrieve");
5841
5842 assert!(
5843 !bundle.uncovered_terms.iter().any(|t| t.starts_with("retr")),
5844 "`retry` must match a corpus that says `retries`; uncovered: {:?}",
5845 bundle.uncovered_terms
5846 );
5847 }
5848
5849 fn ordering_conn() -> rusqlite::Connection {
5856 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
5857 crate::schema::initialize(&conn).expect("init schema");
5858 for (mid, created, text) in [
5859 (
5860 "m_late",
5861 "2026-06-01T09:00:00Z",
5862 "switched the error type to thiserror",
5863 ),
5864 (
5865 "m_early",
5866 "2026-01-15T10:00:00Z",
5867 "ran the thiserror schema migration",
5868 ),
5869 ] {
5870 let normalized = kimetsu_core::memory::normalize_memory_text(text);
5871 conn.execute(
5872 "
5873 INSERT INTO memories (
5874 memory_id, scope, kind, text, normalized_text, confidence,
5875 source_event_id, provenance_snapshot_json, created_at
5876 )
5877 VALUES (?1, 'project', 'fact', ?2, ?3, 1.0, NULL, '{}', ?4)
5878 ",
5879 rusqlite::params![mid, text, normalized, created],
5880 )
5881 .expect("insert memory");
5882 conn.execute(
5883 "INSERT INTO memories_fts (memory_id, text, kind, scope)
5884 VALUES (?1, ?2, 'fact', 'project')",
5885 rusqlite::params![mid, text],
5886 )
5887 .expect("insert fts");
5888 }
5889 conn
5890 }
5891
5892 fn ordering_bundle(conn: &rusqlite::Connection, query: &str) -> ContextBundle {
5893 retrieve_context_with_embedder(
5894 conn,
5895 "/fake-repo",
5896 &kimetsu_core::config::BrokerWeights::default(),
5897 ContextRequest {
5898 stage: "localization".to_string(),
5899 query: query.to_string(),
5900 budget_tokens: 4000,
5901 ..Default::default()
5902 },
5903 &[],
5904 &embeddings::NoopEmbedder,
5905 )
5906 .expect("retrieve")
5907 }
5908
5909 #[test]
5912 fn an_ordering_query_returns_a_dated_chronological_bundle() {
5913 let conn = ordering_conn();
5914 let bundle = ordering_bundle(&conn, "did we run the thiserror migration before or after");
5915
5916 assert!(bundle.chronological, "the query asked about order");
5917 let order: Vec<&str> = bundle
5918 .capsules
5919 .iter()
5920 .filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
5921 .collect();
5922 assert_eq!(order, vec!["m_early", "m_late"], "oldest first");
5923 for (capsule, date) in bundle.capsules.iter().zip(["2026-01-15", "2026-06-01"]) {
5924 assert!(
5925 capsule.summary.contains(&format!("[{date}]")),
5926 "every capsule carries its date; got: {}",
5927 capsule.summary
5928 );
5929 }
5930 }
5931
5932 #[test]
5935 fn an_ordinary_query_is_untouched() {
5936 let conn = ordering_conn();
5937 let bundle = ordering_bundle(&conn, "how do we handle thiserror errors");
5938
5939 assert!(!bundle.chronological);
5940 for capsule in &bundle.capsules {
5941 assert!(
5942 !capsule.summary.contains('['),
5943 "no dates on a non-ordering query; got: {}",
5944 capsule.summary
5945 );
5946 }
5947 }
5948
5949 #[test]
5953 fn ordering_changes_the_rendering_not_the_selection() {
5954 let conn = ordering_conn();
5955 let ordered = ordering_bundle(&conn, "did we run the thiserror migration before or after");
5956 let plain = ordering_bundle(&conn, "did we run the thiserror migration");
5957
5958 let mut got: Vec<&str> = ordered
5959 .capsules
5960 .iter()
5961 .map(|c| c.expansion_handle.as_str())
5962 .collect();
5963 let mut want: Vec<&str> = plain
5964 .capsules
5965 .iter()
5966 .map(|c| c.expansion_handle.as_str())
5967 .collect();
5968 got.sort_unstable();
5969 want.sort_unstable();
5970 assert_eq!(got, want, "same capsules, different order");
5971 }
5972
5973 #[test]
5976 fn the_dates_are_counted_against_the_budget() {
5977 let conn = ordering_conn();
5978 let ordered = ordering_bundle(&conn, "did we run the thiserror migration before or after");
5979 let plain = ordering_bundle(&conn, "did we run the thiserror migration");
5980 assert!(
5981 ordered.used_tokens > plain.used_tokens,
5982 "dated: {} vs plain: {}",
5983 ordered.used_tokens,
5984 plain.used_tokens
5985 );
5986 assert_eq!(
5987 ordered.used_tokens,
5988 ordered
5989 .capsules
5990 .iter()
5991 .map(|c| c.token_estimate)
5992 .sum::<u32>(),
5993 "used_tokens must match what was actually rendered"
5994 );
5995 }
5996}