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 tags: Vec<String>,
287 pub min_score: f32,
292 pub max_capsules: usize,
295 pub prefer_roles: Vec<String>,
299 pub kinds: Vec<String>,
306 pub min_semantic_score: f32,
314 pub min_lexical_coverage: f32,
324 pub task_kind: TaskKind,
329}
330
331#[derive(Debug, Clone)]
332pub struct ContextBundle {
333 pub stage: String,
334 pub budget_tokens: u32,
335 pub used_tokens: u32,
336 pub capsules: Vec<ContextCapsule>,
337 pub excluded: Vec<ContextCapsule>,
338 pub skipped: bool,
341 pub top_score: f32,
344}
345
346#[derive(Debug, Clone)]
352pub(crate) struct Candidate {
353 pub(crate) capsule: ContextCapsule,
354 pub(crate) raw_relevance: f32,
355 pub(crate) embedding: Option<Vec<f32>>,
361 pub(crate) cosine: Option<f32>,
365}
366
367pub fn retrieve_context(
368 conn: &Connection,
369 repo_root: &str,
370 weights: &BrokerWeights,
371 request: ContextRequest,
372) -> KimetsuResult<ContextBundle> {
373 retrieve_context_multi(conn, repo_root, weights, request, &[])
374}
375
376pub fn retrieve_context_multi(
393 conn: &Connection,
394 repo_root: &str,
395 weights: &BrokerWeights,
396 request: ContextRequest,
397 extra_memory_conns: &[&Connection],
398) -> KimetsuResult<ContextBundle> {
399 let embedder = embeddings::open_default_embedder();
400 retrieve_context_with_embedder(
401 conn,
402 repo_root,
403 weights,
404 request,
405 extra_memory_conns,
406 embedder,
407 )
408}
409
410pub fn retrieve_context_with_embedder(
422 conn: &Connection,
423 repo_root: &str,
424 weights: &BrokerWeights,
425 request: ContextRequest,
426 extra_memory_conns: &[&Connection],
427 embedder: &dyn Embedder,
428) -> KimetsuResult<ContextBundle> {
429 retrieve_context_with_embedder_and_backend(
430 conn,
431 repo_root,
432 weights,
433 request,
434 extra_memory_conns,
435 embedder,
436 &crate::backend::FlatBackend,
437 )
438}
439
440pub(crate) fn retrieve_context_with_embedder_and_backend(
453 conn: &Connection,
454 repo_root: &str,
455 weights: &BrokerWeights,
456 request: ContextRequest,
457 extra_memory_conns: &[&Connection],
458 embedder: &dyn Embedder,
459 backend: &dyn crate::backend::RetrievalBackend,
460) -> KimetsuResult<ContextBundle> {
461 let query_embedding = QueryEmbedding::from_embedder(embedder, &request.query);
462 let half_life_days = weights.decay_half_life_days;
463 let mut candidates = Vec::new();
464 candidates.extend(backend.memory_candidates(
465 conn,
466 &request.query,
467 query_embedding.as_ref(),
468 half_life_days,
469 )?);
470 for extra in extra_memory_conns {
471 candidates.extend(backend.memory_candidates(
472 extra,
473 &request.query,
474 query_embedding.as_ref(),
475 half_life_days,
476 )?);
477 }
478 candidates.extend(repo_file_candidates(conn, repo_root, &request.query, 30)?);
479 candidates.extend(manifest_candidates(conn, repo_root, &request.query)?);
480
481 if !request.kinds.is_empty() {
488 candidates.retain(|c| {
489 request
490 .kinds
491 .iter()
492 .any(|k| capsule_matches_kind(&c.capsule, k))
493 });
494 }
495
496 if request.min_lexical_coverage > 0.0 {
513 let content = content_tokens(&request.query);
514 if !content.is_empty() {
515 let idf = corpus_token_idf(conn, &content)?;
516 let total_idf: f32 = content
517 .iter()
518 .map(|t| idf.get(t).copied().unwrap_or(0.0))
519 .sum();
520 if total_idf > f32::EPSILON {
523 candidates.retain(|c| {
524 if c.capsule.kind != "memory" {
525 return true; }
527 if c.cosine.is_some_and(|cos| cos >= SEMANTIC_KEEP_COSINE) {
530 return true;
531 }
532 weighted_coverage(&content, &idf, &c.capsule.summary)
533 >= request.min_lexical_coverage
534 });
535 }
536 }
537 }
538
539 let stage_weights = weights_for_stage(weights, &request.stage);
542 let effective_weights = weights_for_task_kind(stage_weights, request.task_kind);
543 normalize_and_score(&mut candidates, effective_weights);
544
545 let kind_role_hints = task_kind_prefer_roles(request.task_kind);
548 let mut effective_prefer_roles: Vec<String> = request.prefer_roles.clone();
549 for &hint in kind_role_hints {
550 let hint_s = hint.to_string();
551 if !effective_prefer_roles.contains(&hint_s) {
552 effective_prefer_roles.push(hint_s);
553 }
554 }
555
556 if !request.tags.is_empty() || !effective_prefer_roles.is_empty() {
568 let tags_lc: Vec<String> = request
569 .tags
570 .iter()
571 .map(|t| t.to_ascii_lowercase())
572 .collect();
573 for c in &mut candidates {
574 let summary_lc = c.capsule.summary.to_ascii_lowercase();
575 if !tags_lc.is_empty() && tags_lc.iter().any(|t| summary_lc.contains(t.as_str())) {
576 c.capsule.score *= 1.4;
577 }
578 if !effective_prefer_roles.is_empty()
579 && effective_prefer_roles.iter().any(|r| {
580 if c.capsule.kind == "memory" {
588 capsule_matches_kind(&c.capsule, r.as_str())
589 } else {
590 c.capsule.kind.contains(r.as_str())
591 }
592 })
593 {
594 c.capsule.score *= 1.3;
595 }
596 }
597 }
598
599 if query_embedding.is_some() && request.min_semantic_score > 0.0 {
614 candidates.retain(|c| {
615 match c.cosine {
618 Some(cos) => cos >= request.min_semantic_score,
619 None => true,
620 }
621 });
622 }
623
624 candidates.sort_by(|a, b| {
638 b.capsule
639 .score
640 .partial_cmp(&a.capsule.score)
641 .unwrap_or(Ordering::Equal)
642 .then_with(|| {
643 b.capsule
644 .freshness
645 .partial_cmp(&a.capsule.freshness)
646 .unwrap_or(Ordering::Equal)
647 })
648 .then_with(|| a.capsule.expansion_handle.cmp(&b.capsule.expansion_handle))
653 });
654
655 let embedding_mmr_ran = query_embedding.is_some() && !candidates.is_empty();
658 let candidates = if embedding_mmr_ran {
659 apply_candidate_mmr_diversity(candidates, 0.7)
660 } else {
661 candidates
662 };
663
664 let mut capsules = candidates
665 .into_iter()
666 .map(|candidate| candidate.capsule)
667 .collect::<Vec<_>>();
668
669 if !embedding_mmr_ran {
672 capsules.sort_by(|left, right| {
673 right
674 .score
675 .partial_cmp(&left.score)
676 .unwrap_or(Ordering::Equal)
677 .then_with(|| {
678 right
679 .freshness
680 .partial_cmp(&left.freshness)
681 .unwrap_or(Ordering::Equal)
682 })
683 .then_with(|| left.expansion_handle.cmp(&right.expansion_handle))
685 });
686 }
687
688 let top_score = capsules.first().map(|c| c.score).unwrap_or(0.0);
691 if request.min_score > 0.0 && top_score < request.min_score {
692 return Ok(ContextBundle {
693 stage: request.stage,
694 budget_tokens: request.budget_tokens,
695 used_tokens: 0,
696 capsules: Vec::new(),
697 excluded: capsules,
698 skipped: true,
699 top_score,
700 });
701 }
702
703 let capsules = apply_mmr_diversity(capsules, 0.7);
711
712 let capsule_budget = request.budget_tokens / 2;
713 let mut used_tokens = 0u32;
714 let mut included = Vec::new();
715 let mut excluded = Vec::new();
716
717 for capsule in capsules {
718 if request.max_capsules > 0 && included.len() >= request.max_capsules {
720 excluded.push(capsule);
721 continue;
722 }
723 if used_tokens.saturating_add(capsule.token_estimate) <= capsule_budget {
724 used_tokens += capsule.token_estimate;
725 included.push(capsule);
726 } else {
727 excluded.push(capsule);
728 }
729 }
730
731 Ok(ContextBundle {
732 stage: request.stage,
733 budget_tokens: request.budget_tokens,
734 used_tokens,
735 capsules: included,
736 excluded,
737 skipped: false,
738 top_score,
739 })
740}
741
742pub fn search_memories_including_expired(
754 conn: &Connection,
755 limit: u32,
756) -> KimetsuResult<Vec<ContextCapsule>> {
757 let mut stmt = conn.prepare_cached(
758 "
759 SELECT memory_id, scope, kind, text, confidence, created_at,
760 use_count, usefulness_score, valid_from, valid_to
761 FROM memories
762 WHERE invalidated_at IS NULL
763 AND superseded_by IS NULL
764 ORDER BY created_at DESC
765 LIMIT ?1
766 ",
767 )?;
768 let rows = stmt.query_map(params![limit], |row| {
769 Ok((
770 row.get::<_, String>(0)?,
771 row.get::<_, String>(1)?,
772 row.get::<_, String>(2)?,
773 row.get::<_, String>(3)?,
774 row.get::<_, f32>(4)?,
775 row.get::<_, String>(5)?,
776 row.get::<_, i64>(6)?,
777 row.get::<_, f64>(7)?,
778 row.get::<_, Option<String>>(8)?,
779 row.get::<_, Option<String>>(9)?,
780 ))
781 })?;
782 let now_utc = OffsetDateTime::now_utc();
783 let now_rfc3339 = now_utc
784 .format(&time::format_description::well_known::Rfc3339)
785 .unwrap_or_default();
786 let mut capsules = Vec::new();
787 for row in rows {
788 let (
789 memory_id,
790 scope,
791 kind,
792 text,
793 confidence,
794 created_at,
795 _use_count,
796 _usefulness,
797 _valid_from,
798 valid_to,
799 ) = row?;
800 let freshness = freshness(&created_at);
801 let scope_weight = scope_weight(&scope);
802 let suffix = if let Some(ref vt) = valid_to {
804 if vt.as_str() < now_rfc3339.as_str() {
805 format!(" [expired valid_to={vt}]")
806 } else {
807 format!(" [valid_to={vt}]")
808 }
809 } else {
810 String::new()
811 };
812 capsules.push(ContextCapsule {
813 id: new_id().to_string(),
814 kind: "memory".to_string(),
815 summary: format!("{scope}:{kind} - {text}{suffix}"),
816 token_estimate: estimate_tokens(&text) + 8,
817 expansion_handle: format!("memory:{memory_id}"),
818 provenance: vec![ProvenanceRef {
819 source: "Memory".to_string(),
820 id: memory_id,
821 excerpt: Some(excerpt(&text)),
822 }],
823 confidence,
824 freshness,
825 relevance: 0.0,
826 scope_weight,
827 score: 0.0,
828 });
829 }
830 Ok(capsules)
831}
832
833pub fn search_repo_files(
834 conn: &Connection,
835 repo_root: &str,
836 query: &str,
837 limit: u32,
838) -> KimetsuResult<Vec<ContextCapsule>> {
839 let candidates = repo_file_candidates(conn, repo_root, query, limit)?;
840 let mut capsules = candidates
841 .into_iter()
842 .map(|mut candidate| {
843 candidate.capsule.relevance = candidate.raw_relevance;
844 candidate.capsule.score = candidate.raw_relevance;
845 candidate.capsule
846 })
847 .collect::<Vec<_>>();
848 capsules.sort_by(|left, right| {
849 right
850 .score
851 .partial_cmp(&left.score)
852 .unwrap_or(Ordering::Equal)
853 .then_with(|| left.expansion_handle.cmp(&right.expansion_handle))
854 });
855 Ok(capsules)
856}
857
858#[cfg(feature = "embeddings")]
870fn memory_ann_candidates(
871 conn: &Connection,
872 qe: &QueryEmbedding,
873 k: u32,
874 query_tokens: &[String],
875 half_life_days: f32,
876) -> KimetsuResult<Vec<Candidate>> {
877 let handle = crate::ann::handle_for_query(conn, qe.vector.len(), &qe.model_id)?;
879 let hits = handle
880 .read()
881 .unwrap_or_else(|p| p.into_inner())
882 .search(&qe.vector, k as usize)?;
883 let knn_rowids: Vec<i64> = hits.into_iter().map(|(rowid, _dist)| rowid).collect();
887 if knn_rowids.is_empty() {
888 return Ok(Vec::new());
889 }
890
891 let placeholders: String = knn_rowids
893 .iter()
894 .enumerate()
895 .map(|(i, _)| format!("?{}", i + 1))
896 .collect::<Vec<_>>()
897 .join(", ");
898 let sql = format!(
899 "SELECT memory_id, scope, kind, text, confidence, created_at,
900 use_count, usefulness_score, embedding, embedding_model,
901 last_useful_at
902 FROM memories
903 WHERE invalidated_at IS NULL
904 AND superseded_by IS NULL
905 AND (valid_to IS NULL OR valid_to > datetime('now'))
906 AND embedding_model = ?{model_param}
907 AND rowid IN ({placeholders})",
908 model_param = knn_rowids.len() + 1
909 );
910 let mut stmt = conn.prepare(&sql)?;
911 let mut params_vec: Vec<&dyn rusqlite::ToSql> = knn_rowids
912 .iter()
913 .map(|n| n as &dyn rusqlite::ToSql)
914 .collect();
915 params_vec.push(&qe.model_id);
916 let rows_iter = stmt.query_map(params_vec.as_slice(), |row| {
917 Ok((
918 row.get::<_, String>(0)?,
919 row.get::<_, String>(1)?,
920 row.get::<_, String>(2)?,
921 row.get::<_, String>(3)?,
922 row.get::<_, f32>(4)?,
923 row.get::<_, String>(5)?,
924 row.get::<_, i64>(6)?,
925 row.get::<_, f64>(7)?,
926 row.get::<_, Option<Vec<u8>>>(8)?,
927 row.get::<_, Option<String>>(9)?,
928 row.get::<_, Option<String>>(10)?,
929 ))
930 })?;
931
932 let mut candidates = Vec::new();
933 for row in rows_iter {
934 let (
935 memory_id,
936 scope,
937 kind,
938 text,
939 confidence,
940 created_at,
941 use_count,
942 usefulness_score,
943 embedding,
944 embedding_model,
945 last_useful_at,
946 ) = row?;
947 let (cosine, row_vec) =
948 compute_cosine_and_vec(Some(qe), embedding.as_deref(), embedding_model.as_deref());
949 if let Some(candidate) = memory_row_to_candidate(
950 query_tokens,
951 memory_id,
952 scope,
953 kind,
954 text,
955 confidence,
956 created_at,
957 use_count,
958 usefulness_score,
959 last_useful_at,
960 half_life_days,
961 None, cosine,
963 row_vec,
964 ) {
965 candidates.push(candidate);
966 }
967 }
968 Ok(candidates)
969}
970
971pub(crate) fn memory_candidates_flat(
977 conn: &Connection,
978 query: &str,
979 query_embedding: Option<&QueryEmbedding>,
980 half_life_days: f32,
981) -> KimetsuResult<Vec<Candidate>> {
982 memory_candidates(conn, query, query_embedding, half_life_days)
983}
984
985fn memory_candidates(
986 conn: &Connection,
987 query: &str,
988 query_embedding: Option<&QueryEmbedding>,
989 half_life_days: f32,
990) -> KimetsuResult<Vec<Candidate>> {
991 let query_tokens = query_tokens(query);
992
993 #[cfg(feature = "embeddings")]
998 if let Some(qe) = query_embedding {
999 let fts_candidates = if let Some(fts_query) = fts_query(query) {
1001 memory_fts_candidates(
1002 conn,
1003 &query_tokens,
1004 &fts_query,
1005 80,
1006 Some(qe),
1007 half_life_days,
1008 )?
1009 } else {
1010 Vec::new()
1011 };
1012
1013 let ann_candidates = memory_ann_candidates(conn, qe, 80, &query_tokens, half_life_days)?;
1015
1016 let mut seen: HashMap<String, usize> = HashMap::new();
1021 let mut merged: Vec<Candidate> = Vec::new();
1022
1023 for candidate in fts_candidates.into_iter().chain(ann_candidates) {
1024 let mid = candidate
1026 .capsule
1027 .expansion_handle
1028 .strip_prefix("memory:")
1029 .unwrap_or(&candidate.capsule.expansion_handle)
1030 .to_string();
1031 if let Some(&idx) = seen.get(&mid) {
1032 if candidate.raw_relevance > merged[idx].raw_relevance {
1034 merged[idx] = candidate;
1035 }
1036 } else {
1037 seen.insert(mid, merged.len());
1038 merged.push(candidate);
1039 }
1040 }
1041
1042 return Ok(merged);
1043 }
1044
1045 if let Some(fts_query) = fts_query(query) {
1047 let candidates = memory_fts_candidates(
1048 conn,
1049 &query_tokens,
1050 &fts_query,
1051 80,
1052 query_embedding,
1053 half_life_days,
1054 )?;
1055 if !candidates.is_empty() {
1056 return Ok(candidates);
1057 }
1058 }
1059
1060 latest_memory_candidates(conn, &query_tokens, 200, query_embedding, half_life_days)
1061}
1062
1063fn latest_memory_candidates(
1064 conn: &Connection,
1065 query_tokens: &[String],
1066 limit: u32,
1067 query_embedding: Option<&QueryEmbedding>,
1068 half_life_days: f32,
1069) -> KimetsuResult<Vec<Candidate>> {
1070 let mut stmt = conn.prepare_cached(
1081 "
1082 SELECT memory_id, scope, kind, text, confidence, created_at,
1083 use_count, usefulness_score, embedding, embedding_model,
1084 last_useful_at
1085 FROM memories
1086 WHERE invalidated_at IS NULL
1087 AND superseded_by IS NULL
1088 AND (valid_to IS NULL OR valid_to > datetime('now'))
1089 ORDER BY created_at DESC
1090 LIMIT ?1
1091 ",
1092 )?;
1093
1094 let rows = stmt.query_map(params![limit], |row| {
1095 Ok((
1096 row.get::<_, String>(0)?,
1097 row.get::<_, String>(1)?,
1098 row.get::<_, String>(2)?,
1099 row.get::<_, String>(3)?,
1100 row.get::<_, f32>(4)?,
1101 row.get::<_, String>(5)?,
1102 row.get::<_, i64>(6)?,
1103 row.get::<_, f64>(7)?,
1104 row.get::<_, Option<Vec<u8>>>(8)?,
1105 row.get::<_, Option<String>>(9)?,
1106 row.get::<_, Option<String>>(10)?,
1107 ))
1108 })?;
1109
1110 let mut candidates = Vec::new();
1111 for row in rows {
1112 let (
1113 memory_id,
1114 scope,
1115 kind,
1116 text,
1117 confidence,
1118 created_at,
1119 use_count,
1120 usefulness_score,
1121 embedding,
1122 embedding_model,
1123 last_useful_at,
1124 ) = row?;
1125 let (cosine, row_vec) = compute_cosine_and_vec(
1126 query_embedding,
1127 embedding.as_deref(),
1128 embedding_model.as_deref(),
1129 );
1130 if let Some(candidate) = memory_row_to_candidate(
1131 query_tokens,
1132 memory_id,
1133 scope,
1134 kind,
1135 text,
1136 confidence,
1137 created_at,
1138 use_count,
1139 usefulness_score,
1140 last_useful_at,
1141 half_life_days,
1142 None,
1143 cosine,
1144 row_vec,
1145 ) {
1146 candidates.push(candidate);
1147 }
1148 }
1149 Ok(candidates)
1150}
1151
1152fn memory_fts_candidates(
1153 conn: &Connection,
1154 query_tokens: &[String],
1155 fts_query: &str,
1156 limit: u32,
1157 query_embedding: Option<&QueryEmbedding>,
1158 half_life_days: f32,
1159) -> KimetsuResult<Vec<Candidate>> {
1160 let mut stmt = conn.prepare_cached(
1161 "
1162 SELECT m.memory_id, m.scope, m.kind, m.text, m.confidence, m.created_at,
1163 m.use_count, m.usefulness_score, bm25(memories_fts) AS rank,
1164 m.embedding, m.embedding_model, m.last_useful_at
1165 FROM memories_fts
1166 JOIN memories m
1167 ON m.memory_id = memories_fts.memory_id
1168 WHERE m.invalidated_at IS NULL
1169 AND m.superseded_by IS NULL
1170 AND (m.valid_to IS NULL OR m.valid_to > datetime('now'))
1171 AND memories_fts MATCH ?1
1172 ORDER BY rank
1173 LIMIT ?2
1174 ",
1175 )?;
1176
1177 let rows = stmt.query_map(params![fts_query, limit], |row| {
1178 Ok((
1179 row.get::<_, String>(0)?,
1180 row.get::<_, String>(1)?,
1181 row.get::<_, String>(2)?,
1182 row.get::<_, String>(3)?,
1183 row.get::<_, f32>(4)?,
1184 row.get::<_, String>(5)?,
1185 row.get::<_, i64>(6)?,
1186 row.get::<_, f64>(7)?,
1187 row.get::<_, f64>(8)?,
1188 row.get::<_, Option<Vec<u8>>>(9)?,
1189 row.get::<_, Option<String>>(10)?,
1190 row.get::<_, Option<String>>(11)?,
1191 ))
1192 })?;
1193
1194 let mut candidates = Vec::new();
1195 for row in rows {
1196 let (
1197 memory_id,
1198 scope,
1199 kind,
1200 text,
1201 confidence,
1202 created_at,
1203 use_count,
1204 usefulness_score,
1205 rank,
1206 embedding,
1207 embedding_model,
1208 last_useful_at,
1209 ) = row?;
1210 let fts_relevance = (-rank as f32).max(0.0);
1211 let (cosine, row_vec) = compute_cosine_and_vec(
1212 query_embedding,
1213 embedding.as_deref(),
1214 embedding_model.as_deref(),
1215 );
1216 if let Some(candidate) = memory_row_to_candidate(
1217 query_tokens,
1218 memory_id,
1219 scope,
1220 kind,
1221 text,
1222 confidence,
1223 created_at,
1224 use_count,
1225 usefulness_score,
1226 last_useful_at,
1227 half_life_days,
1228 Some(fts_relevance),
1229 cosine,
1230 row_vec,
1231 ) {
1232 candidates.push(candidate);
1233 }
1234 }
1235 Ok(candidates)
1236}
1237
1238fn compute_cosine_and_vec(
1262 query_embedding: Option<&QueryEmbedding>,
1263 row_bytes: Option<&[u8]>,
1264 row_model: Option<&str>,
1265) -> (Option<f32>, Option<Vec<f32>>) {
1266 let q = match query_embedding {
1267 Some(q) => q,
1268 None => return (None, None),
1269 };
1270 let bytes = match row_bytes {
1271 Some(b) => b,
1272 None => return (None, None),
1273 };
1274 let model = match row_model {
1275 Some(m) => m,
1276 None => return (None, None),
1277 };
1278 if model != q.model_id {
1279 return (None, None);
1280 }
1281 let row_vec = match decode_embedding(bytes, Some(q.vector.len())) {
1282 Ok(v) => v,
1283 Err(_) => return (None, None),
1284 };
1285 let score = cosine_similarity(&q.vector, &row_vec);
1286 (Some(score), Some(row_vec))
1287}
1288
1289#[allow(clippy::too_many_arguments)]
1290fn memory_row_to_candidate(
1291 query_tokens: &[String],
1292 memory_id: String,
1293 scope: String,
1294 kind: String,
1295 text: String,
1296 confidence: f32,
1297 created_at: String,
1298 use_count: i64,
1299 usefulness_score: f64,
1300 last_useful_at: Option<String>,
1301 half_life_days: f32,
1302 raw_relevance_override: Option<f32>,
1303 cosine_score: Option<f32>,
1304 row_embedding: Option<Vec<f32>>,
1309) -> Option<Candidate> {
1310 let lexical = lexical_relevance(query_tokens, &format!("{kind} {text}"));
1311 let lexical_term = raw_relevance_override.unwrap_or(lexical).max(lexical);
1312
1313 let raw_relevance = match cosine_score {
1325 Some(c) => {
1326 let normalized_cos = ((c + 1.0) * 0.5).clamp(0.0, 1.0);
1327 (1.0 - DEFAULT_HYBRID_ALPHA) * lexical_term + DEFAULT_HYBRID_ALPHA * normalized_cos
1328 }
1329 None => lexical_term,
1330 };
1331
1332 if raw_relevance <= 0.0 && !query_tokens.is_empty() {
1338 return None;
1339 }
1340
1341 let freshness = freshness(&created_at);
1342 let scope_weight = scope_weight(&scope);
1343 let raw_multiplier = usefulness_multiplier(usefulness_score as f32, use_count as u32);
1349 let decay = usefulness_decay(last_useful_at.as_deref(), &created_at, half_life_days);
1350 let multiplier = 1.0 + (raw_multiplier - 1.0) * decay;
1351 let biased_relevance = raw_relevance * multiplier;
1352 Some(Candidate {
1353 raw_relevance: biased_relevance,
1354 embedding: row_embedding,
1355 cosine: cosine_score,
1356 capsule: ContextCapsule {
1357 id: new_id().to_string(),
1358 kind: "memory".to_string(),
1359 summary: format!("{scope}:{kind} - {text}"),
1360 token_estimate: estimate_tokens(&text) + 8,
1361 expansion_handle: format!("memory:{memory_id}"),
1362 provenance: vec![ProvenanceRef {
1363 source: "Memory".to_string(),
1364 id: memory_id,
1365 excerpt: Some(excerpt(&text)),
1366 }],
1367 confidence,
1368 freshness,
1369 relevance: 0.0,
1370 scope_weight,
1371 score: 0.0,
1372 },
1373 })
1374}
1375
1376pub(crate) fn usefulness_decay(
1400 last_useful_at: Option<&str>,
1401 created_at: &str,
1402 half_life_days: f32,
1403) -> f32 {
1404 if half_life_days <= 0.0 {
1405 return 1.0;
1406 }
1407 let reference = last_useful_at.unwrap_or(created_at);
1408 let Ok(reference_ts) =
1409 OffsetDateTime::parse(reference, &time::format_description::well_known::Rfc3339)
1410 else {
1411 return 1.0;
1412 };
1413 let age = OffsetDateTime::now_utc() - reference_ts;
1414 let age_days = (age.whole_seconds().max(0) as f32) / 86_400.0;
1415 let exponent = -std::f32::consts::LN_2 * age_days / half_life_days;
1416 exponent.exp().clamp(0.0, 1.0)
1417}
1418
1419pub(crate) fn usefulness_multiplier(usefulness_score: f32, use_count: u32) -> f32 {
1424 const FULL_CONFIDENCE_USES: u32 = 3;
1432 const MULTIPLIER_MIN: f32 = 0.5;
1433 const MULTIPLIER_MAX: f32 = 1.5;
1434 if use_count == 0 {
1435 return 1.0;
1436 }
1437 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);
1440 let confidence = (use_count as f32 / FULL_CONFIDENCE_USES as f32).min(1.0);
1441 1.0 * (1.0 - confidence) + full_multiplier * confidence
1442}
1443
1444fn repo_file_candidates(
1445 conn: &Connection,
1446 repo_root: &str,
1447 query: &str,
1448 limit: u32,
1449) -> KimetsuResult<Vec<Candidate>> {
1450 let Some(fts_query) = fts_query(query) else {
1451 return Ok(Vec::new());
1452 };
1453
1454 let mut stmt = conn.prepare_cached(
1455 "
1456 SELECT path, snippet, language_guess, bm25(repo_files_fts) AS rank
1457 FROM repo_files_fts
1458 WHERE repo_root = ?1 AND repo_files_fts MATCH ?2
1459 ORDER BY rank
1460 LIMIT ?3
1461 ",
1462 )?;
1463
1464 let rows = stmt.query_map(params![repo_root, fts_query, limit], |row| {
1465 Ok((
1466 row.get::<_, String>(0)?,
1467 row.get::<_, String>(1)?,
1468 row.get::<_, String>(2)?,
1469 row.get::<_, f64>(3)?,
1470 ))
1471 })?;
1472
1473 let mut candidates = Vec::new();
1474 for row in rows {
1475 let (path, snippet, language, rank) = row?;
1476 let raw_relevance = (-rank as f32).max(0.0);
1477 let summary = format!("{path} ({language}) - {}", excerpt(&snippet));
1478 let token_estimate = estimate_tokens(&summary) + 8;
1479 candidates.push(Candidate {
1480 raw_relevance,
1481 embedding: None,
1482 cosine: None,
1483 capsule: ContextCapsule {
1484 id: new_id().to_string(),
1485 kind: "repo_file".to_string(),
1486 summary,
1487 token_estimate,
1488 expansion_handle: format!("file:{path}"),
1489 provenance: vec![ProvenanceRef {
1490 source: "RepoFile".to_string(),
1491 id: path.clone(),
1492 excerpt: Some(excerpt(&snippet)),
1493 }],
1494 confidence: 0.9,
1495 freshness: 1.0,
1496 relevance: 0.0,
1497 scope_weight: 0.9,
1498 score: 0.0,
1499 },
1500 });
1501 }
1502 Ok(candidates)
1503}
1504
1505fn manifest_candidates(
1506 conn: &Connection,
1507 repo_root: &str,
1508 query: &str,
1509) -> KimetsuResult<Vec<Candidate>> {
1510 if let Some(fts_query) = fts_query(query) {
1511 let candidates = manifest_fts_candidates(conn, repo_root, &fts_query, 30)?;
1512 if !candidates.is_empty() {
1513 return Ok(candidates);
1514 }
1515 }
1516
1517 let query_tokens = query_tokens(query);
1518 let mut stmt = conn.prepare_cached(
1519 "
1520 SELECT manifest_path, manifest_kind, parsed_summary_json
1521 FROM repo_manifests
1522 WHERE repo_root = ?1
1523 ORDER BY manifest_path
1524 ",
1525 )?;
1526
1527 let rows = stmt.query_map(params![repo_root], |row| {
1528 Ok((
1529 row.get::<_, String>(0)?,
1530 row.get::<_, String>(1)?,
1531 row.get::<_, String>(2)?,
1532 ))
1533 })?;
1534
1535 let mut candidates = Vec::new();
1536 for row in rows {
1537 let (path, kind, summary_json) = row?;
1538 let raw_relevance =
1539 lexical_relevance(&query_tokens, &format!("{path} {kind} {summary_json}"));
1540 if raw_relevance <= 0.0 && !query_tokens.is_empty() {
1541 continue;
1542 }
1543 let summary = format!("{path} manifest ({kind})");
1544 let token_estimate = estimate_tokens(&summary) + 8;
1545 candidates.push(Candidate {
1546 raw_relevance,
1547 embedding: None,
1548 cosine: None,
1549 capsule: ContextCapsule {
1550 id: new_id().to_string(),
1551 kind: "repo_manifest".to_string(),
1552 summary,
1553 token_estimate,
1554 expansion_handle: format!("file:{path}"),
1555 provenance: vec![ProvenanceRef {
1556 source: "Manifest".to_string(),
1557 id: path,
1558 excerpt: Some(excerpt(&summary_json)),
1559 }],
1560 confidence: 0.95,
1561 freshness: 1.0,
1562 relevance: 0.0,
1563 scope_weight: 0.9,
1564 score: 0.0,
1565 },
1566 });
1567 }
1568 Ok(candidates)
1569}
1570
1571fn manifest_fts_candidates(
1572 conn: &Connection,
1573 repo_root: &str,
1574 fts_query: &str,
1575 limit: u32,
1576) -> KimetsuResult<Vec<Candidate>> {
1577 let mut stmt = conn.prepare_cached(
1578 "
1579 SELECT manifest_path, manifest_kind, parsed_summary_json,
1580 bm25(repo_manifests_fts) AS rank
1581 FROM repo_manifests_fts
1582 WHERE repo_root = ?1 AND repo_manifests_fts MATCH ?2
1583 ORDER BY rank
1584 LIMIT ?3
1585 ",
1586 )?;
1587
1588 let rows = stmt.query_map(params![repo_root, fts_query, limit], |row| {
1589 Ok((
1590 row.get::<_, String>(0)?,
1591 row.get::<_, String>(1)?,
1592 row.get::<_, String>(2)?,
1593 row.get::<_, f64>(3)?,
1594 ))
1595 })?;
1596
1597 let mut candidates = Vec::new();
1598 for row in rows {
1599 let (path, kind, summary_json, rank) = row?;
1600 let raw_relevance = (-rank as f32).max(0.0);
1601 let summary = format!("{path} manifest ({kind})");
1602 let token_estimate = estimate_tokens(&summary) + 8;
1603 candidates.push(Candidate {
1604 raw_relevance,
1605 embedding: None,
1606 cosine: None,
1607 capsule: ContextCapsule {
1608 id: new_id().to_string(),
1609 kind: "repo_manifest".to_string(),
1610 summary,
1611 token_estimate,
1612 expansion_handle: format!("file:{path}"),
1613 provenance: vec![ProvenanceRef {
1614 source: "Manifest".to_string(),
1615 id: path,
1616 excerpt: Some(excerpt(&summary_json)),
1617 }],
1618 confidence: 0.95,
1619 freshness: 1.0,
1620 relevance: 0.0,
1621 scope_weight: 0.9,
1622 score: 0.0,
1623 },
1624 });
1625 }
1626 Ok(candidates)
1627}
1628
1629fn normalize_and_score(candidates: &mut [Candidate], weights: StageWeights) {
1630 let mut max_by_kind = HashMap::<String, f32>::new();
1631 for candidate in candidates.iter() {
1632 max_by_kind
1633 .entry(candidate.capsule.kind.clone())
1634 .and_modify(|max| *max = (*max).max(candidate.raw_relevance))
1635 .or_insert(candidate.raw_relevance);
1636 }
1637
1638 for candidate in candidates {
1639 let max = max_by_kind
1640 .get(&candidate.capsule.kind)
1641 .copied()
1642 .unwrap_or(0.0);
1643 let relevance = if max <= f32::EPSILON {
1644 if candidate.raw_relevance > 0.0 {
1645 1.0
1646 } else {
1647 0.0
1648 }
1649 } else {
1650 (candidate.raw_relevance / max).clamp(0.0, 1.0)
1651 };
1652 candidate.capsule.relevance = relevance;
1653 candidate.capsule.score = weights.relevance * relevance
1654 + weights.confidence * candidate.capsule.confidence
1655 + weights.freshness * candidate.capsule.freshness
1656 + weights.scope * candidate.capsule.scope_weight;
1657 }
1658}
1659
1660fn weights_for_stage(weights: &BrokerWeights, stage: &str) -> StageWeights {
1661 match stage {
1662 "localization" => weights.localization.clone(),
1663 "patch_plan" => weights.patch_plan.clone(),
1664 "verification" => weights.verification.clone(),
1665 "review" => weights.review.clone(),
1666 _ => None,
1667 }
1668 .unwrap_or(StageWeights {
1669 relevance: weights.relevance,
1670 confidence: weights.confidence,
1671 freshness: weights.freshness,
1672 scope: weights.scope,
1673 })
1674}
1675
1676pub(crate) fn scope_weight_pub(scope: &str) -> f32 {
1679 scope_weight(scope)
1680}
1681
1682fn scope_weight(scope: &str) -> f32 {
1683 match scope.parse::<MemoryScope>() {
1684 Ok(MemoryScope::Run) => 1.0,
1685 Ok(MemoryScope::Repo) => 0.9,
1686 Ok(MemoryScope::Project) => 0.7,
1687 Ok(MemoryScope::GlobalUser) => 0.5,
1688 Err(_) => 0.3,
1689 }
1690}
1691
1692pub(crate) fn freshness_pub(created_at: &str) -> f32 {
1695 freshness(created_at)
1696}
1697
1698fn freshness(created_at: &str) -> f32 {
1699 let Ok(created_at) =
1700 OffsetDateTime::parse(created_at, &time::format_description::well_known::Rfc3339)
1701 else {
1702 return 0.5;
1703 };
1704 let age = OffsetDateTime::now_utc() - created_at;
1705 let age_days = age.whole_seconds().max(0) as f32 / 86_400.0;
1706 (-age_days / 30.0).exp().clamp(0.0, 1.0)
1707}
1708
1709const SEMANTIC_KEEP_COSINE: f32 = 0.20;
1714
1715const STOPWORDS: &[&str] = &[
1720 "the", "and", "for", "are", "but", "not", "you", "your", "with", "this", "that", "these",
1721 "those", "from", "into", "about", "what", "whats", "which", "who", "whom", "how", "why",
1722 "when", "where", "can", "could", "would", "should", "will", "shall", "does", "did", "was",
1723 "were", "been", "being", "have", "has", "had", "its", "it", "is", "as", "at", "by", "of", "to",
1724 "in", "on", "or", "an", "be", "do", "me", "my", "we", "us", "our", "im", "ive", "let", "lets",
1725 "please", "tell", "give", "show", "want", "need", "get", "got", "use", "using", "there",
1726 "their", "they", "them", "then", "than", "some", "any", "all", "more", "most", "such", "via",
1727 "per",
1728];
1729
1730fn content_tokens(query: &str) -> Vec<String> {
1735 let mut seen = std::collections::HashSet::new();
1736 query
1737 .split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
1738 .map(str::trim)
1739 .filter(|part| part.len() >= 2)
1740 .map(str::to_ascii_lowercase)
1741 .filter(|t| !STOPWORDS.contains(&t.as_str()))
1742 .map(|t| light_stem(&t).to_string())
1745 .filter(|t| seen.insert(t.clone()))
1746 .collect()
1747}
1748
1749fn corpus_token_idf(conn: &Connection, tokens: &[String]) -> KimetsuResult<HashMap<String, f32>> {
1767 let mut idf = HashMap::new();
1768 let n: i64 = conn
1769 .query_row(
1770 "SELECT COUNT(*) FROM memories WHERE invalidated_at IS NULL",
1771 [],
1772 |row| row.get(0),
1773 )
1774 .unwrap_or(0);
1775 if n == 0 {
1776 return Ok(idf);
1777 }
1778 let mut stmt = conn.prepare_cached(
1779 "SELECT COUNT(*) FROM memories \
1780 WHERE invalidated_at IS NULL AND lower(text) LIKE ?1 ESCAPE '\\'",
1781 )?;
1782 for token in tokens {
1783 let pattern = format!("%{}%", escape_like(token));
1784 let df: i64 = stmt
1785 .query_row(params![pattern], |row| row.get(0))
1786 .unwrap_or(0);
1787 let weight = if df == 0 {
1789 0.0
1790 } else {
1791 (((n + 1) as f32) / ((df + 1) as f32)).ln().max(0.0)
1792 };
1793 idf.insert(token.clone(), weight);
1794 }
1795 Ok(idf)
1796}
1797
1798fn escape_like(token: &str) -> String {
1801 token
1802 .replace('\\', "\\\\")
1803 .replace('%', "\\%")
1804 .replace('_', "\\_")
1805}
1806
1807fn weighted_coverage(content: &[String], idf: &HashMap<String, f32>, summary: &str) -> f32 {
1814 let haystack = summary.to_ascii_lowercase();
1815 let mut total = 0.0f32;
1816 let mut hit = 0.0f32;
1817 for token in content {
1818 let weight = idf.get(token).copied().unwrap_or(0.0);
1819 total += weight;
1820 if weight > 0.0 && haystack.contains(token.as_str()) {
1821 hit += weight;
1822 }
1823 }
1824 if total <= f32::EPSILON {
1825 0.0
1826 } else {
1827 (hit / total).clamp(0.0, 1.0)
1828 }
1829}
1830
1831fn light_stem(token: &str) -> &str {
1842 for suffix in ["ing", "ed", "es", "s"] {
1843 if let Some(stem) = token.strip_suffix(suffix)
1844 && stem.len() >= 4
1845 {
1846 return stem;
1847 }
1848 }
1849 token
1850}
1851
1852fn query_tokens(query: &str) -> Vec<String> {
1853 let mut tokens: Vec<String> = query
1854 .split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
1855 .map(str::trim)
1856 .filter(|part| part.len() >= 2)
1857 .map(str::to_ascii_lowercase)
1858 .map(|t| light_stem(&t).to_string())
1859 .collect();
1860 let lower = query.to_ascii_lowercase();
1867 for (triggers, expansions) in CLASS_HINTS.iter() {
1868 if triggers.iter().any(|t| lower.contains(t)) {
1869 tokens.extend(expansions.iter().map(|e| e.to_string()));
1870 }
1871 }
1872 tokens
1873}
1874
1875const CLASS_HINTS: &[(&[&str], &[&str])] = &[
1883 (
1884 &[
1885 "build",
1886 "compile",
1887 "make",
1888 "cargo",
1889 "cmake",
1890 "configure",
1891 "install",
1892 "train",
1893 "benchmark",
1894 "test suite",
1895 "ray trace",
1896 "render",
1897 ],
1898 &[
1899 "shell_background",
1900 "shell_status",
1901 "shell_output",
1902 "shell_stop",
1903 "long_running",
1904 ],
1905 ),
1906 (
1907 &[
1908 "edit", "modify", "change", "fix", "update", "patch", "refactor", "rename",
1909 ],
1910 &["edit_file", "apply_patch", "old_string", "new_string"],
1911 ),
1912 (
1913 &[
1914 "read", "inspect", "review", "analyze", "examine", "view", "show",
1915 ],
1916 &["read_file", "offset", "limit", "multi_read"],
1917 ),
1918 (
1919 &["find", "locate", "search", "look up", "discover", "list"],
1920 &["glob", "search_files", "list_files"],
1921 ),
1922 (
1923 &["plan", "step", "checklist", "todo", "task list", "phase"],
1924 &["plan", "todos"],
1925 ),
1926 (
1927 &[
1928 "verify",
1929 "check",
1930 "ensure",
1931 "validate",
1932 "pass test",
1933 "verifier",
1934 ],
1935 &["finish", "verifier", "verification"],
1936 ),
1937 (
1938 &[
1939 "image",
1940 "png",
1941 "jpeg",
1942 "jpg",
1943 "pdf",
1944 "diagram",
1945 "screenshot",
1946 ],
1947 &["view_image", "base64", "sha256"],
1948 ),
1949 (&["delete", "remove", "rm "], &["delete_file", "recursive"]),
1950 (&["rename", "move file", "mv "], &["move_file"]),
1951];
1952
1953fn capsule_matches_kind(capsule: &ContextCapsule, wanted: &str) -> bool {
1958 if capsule.kind == wanted {
1959 return true;
1960 }
1961 if capsule.kind == "memory"
1962 && let Some((prefix, _)) = capsule.summary.split_once(" - ")
1963 && let Some((_scope, mkind)) = prefix.split_once(':')
1964 {
1965 return mkind == wanted;
1966 }
1967 false
1968}
1969
1970pub(crate) fn fts_query(query: &str) -> Option<String> {
1971 let tokens = query_tokens(query);
1972 if tokens.is_empty() {
1973 return None;
1974 }
1975 Some(
1976 tokens
1977 .into_iter()
1978 .take(12)
1979 .map(|token| format!("{token}*"))
1980 .collect::<Vec<_>>()
1981 .join(" OR "),
1982 )
1983}
1984
1985fn apply_candidate_mmr_diversity(mut sorted: Vec<Candidate>, lambda: f32) -> Vec<Candidate> {
2005 if sorted.len() <= 1 {
2006 return sorted;
2007 }
2008 let summaries: Vec<std::collections::HashSet<String>> = sorted
2010 .iter()
2011 .map(|c| summary_token_set(&c.capsule.summary))
2012 .collect();
2013
2014 let mut picked_indices: Vec<usize> = Vec::with_capacity(sorted.len());
2015 let mut remaining: Vec<usize> = (0..sorted.len()).collect();
2016
2017 picked_indices.push(remaining.remove(0));
2019
2020 while !remaining.is_empty() {
2021 let mut best_idx_in_remaining = 0;
2022 let mut best_score = f32::MIN;
2023
2024 for (i, &cand) in remaining.iter().enumerate() {
2025 let mut max_overlap = 0.0f32;
2026 for &p in &picked_indices {
2027 let same_kind = sorted[cand].capsule.kind == sorted[p].capsule.kind;
2030 let raw_overlap = candidate_pair_overlap(
2031 &sorted[cand],
2032 &sorted[p],
2033 &summaries[cand],
2034 &summaries[p],
2035 );
2036 let overlap = if same_kind {
2037 raw_overlap
2038 } else {
2039 raw_overlap * 0.5
2040 };
2041 if overlap > max_overlap {
2042 max_overlap = overlap;
2043 }
2044 }
2045 let mmr = lambda * sorted[cand].capsule.score - (1.0 - lambda) * max_overlap;
2046 if mmr > best_score {
2047 best_score = mmr;
2048 best_idx_in_remaining = i;
2049 }
2050 }
2051 picked_indices.push(remaining.remove(best_idx_in_remaining));
2052 }
2053
2054 let mut taken: Vec<Option<Candidate>> = sorted.drain(..).map(Some).collect();
2056 let mut out = Vec::with_capacity(taken.len());
2057 for idx in picked_indices {
2058 if let Some(c) = taken[idx].take() {
2059 out.push(c);
2060 }
2061 }
2062 out
2063}
2064
2065fn candidate_pair_overlap(
2072 a: &Candidate,
2073 b: &Candidate,
2074 tokens_a: &std::collections::HashSet<String>,
2075 tokens_b: &std::collections::HashSet<String>,
2076) -> f32 {
2077 if let (Some(va), Some(vb)) = (a.embedding.as_deref(), b.embedding.as_deref()) {
2078 cosine_similarity(va, vb).max(0.0)
2083 } else {
2084 jaccard(tokens_a, tokens_b)
2085 }
2086}
2087
2088fn apply_mmr_diversity(mut sorted: Vec<ContextCapsule>, lambda: f32) -> Vec<ContextCapsule> {
2100 if sorted.len() <= 1 {
2101 return sorted;
2102 }
2103 let summaries: Vec<std::collections::HashSet<String>> = sorted
2105 .iter()
2106 .map(|c| summary_token_set(&c.summary))
2107 .collect();
2108 let mut picked_indices: Vec<usize> = Vec::with_capacity(sorted.len());
2109 let mut remaining: Vec<usize> = (0..sorted.len()).collect();
2110
2111 picked_indices.push(remaining.remove(0));
2113
2114 while !remaining.is_empty() {
2115 let mut best_idx_in_remaining = 0;
2116 let mut best_score = f32::MIN;
2117 for (i, &cand) in remaining.iter().enumerate() {
2118 let mut max_overlap = 0.0f32;
2119 for &p in &picked_indices {
2120 let raw = jaccard(&summaries[cand], &summaries[p]);
2121 let overlap = if sorted[cand].kind == sorted[p].kind {
2122 raw
2123 } else {
2124 raw * 0.5
2127 };
2128 if overlap > max_overlap {
2129 max_overlap = overlap;
2130 }
2131 }
2132 let mmr = lambda * sorted[cand].score - (1.0 - lambda) * max_overlap;
2133 if mmr > best_score {
2134 best_score = mmr;
2135 best_idx_in_remaining = i;
2136 }
2137 }
2138 picked_indices.push(remaining.remove(best_idx_in_remaining));
2139 }
2140 let mut out = Vec::with_capacity(sorted.len());
2142 let mut taken: Vec<Option<ContextCapsule>> = sorted.drain(..).map(Some).collect();
2144 for idx in picked_indices {
2145 if let Some(c) = taken[idx].take() {
2146 out.push(c);
2147 }
2148 }
2149 out
2150}
2151
2152fn summary_token_set(s: &str) -> std::collections::HashSet<String> {
2153 s.split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
2154 .filter(|t| t.len() >= 3)
2155 .map(str::to_ascii_lowercase)
2156 .collect()
2157}
2158
2159fn jaccard(a: &std::collections::HashSet<String>, b: &std::collections::HashSet<String>) -> f32 {
2160 if a.is_empty() && b.is_empty() {
2161 return 0.0;
2162 }
2163 let intersection = a.intersection(b).count();
2164 let union = a.union(b).count();
2165 intersection as f32 / union.max(1) as f32
2166}
2167
2168fn lexical_relevance(tokens: &[String], haystack: &str) -> f32 {
2169 if tokens.is_empty() {
2170 return 0.0;
2171 }
2172 let haystack = haystack.to_ascii_lowercase();
2173 let matches = tokens
2174 .iter()
2175 .filter(|token| haystack.contains(token.as_str()))
2176 .count();
2177 matches as f32 / tokens.len() as f32
2178}
2179
2180pub fn estimate_tokens(text: &str) -> u32 {
2181 ((text.split_whitespace().count() as f32) * 1.33).ceil() as u32
2182}
2183
2184pub fn compress_for_render(summary: &str, max_sentences: usize) -> String {
2207 if max_sentences == 0 {
2208 return summary.to_string();
2209 }
2210
2211 let text = if let Some(rest) = summary.strip_prefix('[') {
2213 if let Some(idx) = rest.find(']') {
2215 rest[idx + 1..].trim_start()
2216 } else {
2217 summary
2218 }
2219 } else {
2220 summary
2221 };
2222
2223 let text = if let Some(idx) = text.rfind('(') {
2225 let candidate = text[..idx].trim_end();
2226 let inner = &text[idx + 1..];
2229 if inner.contains(':') && inner.trim_end().ends_with(')') {
2230 candidate
2231 } else {
2232 text
2233 }
2234 } else {
2235 text
2236 };
2237
2238 let (scope_prefix, body) = if let Some(dash_pos) = text.find(" - ") {
2240 let prefix_candidate = &text[..dash_pos];
2241 if !prefix_candidate.contains(' ') && prefix_candidate.contains(':') {
2243 let body_start = dash_pos + 3; (&text[..body_start], &text[body_start..])
2245 } else {
2246 ("", text)
2247 }
2248 } else {
2249 ("", text)
2250 };
2251
2252 let compressed_body = cap_sentences(body, max_sentences);
2254
2255 let result = if scope_prefix.is_empty() {
2257 compressed_body.to_string()
2258 } else {
2259 format!("{scope_prefix}{compressed_body}")
2260 };
2261
2262 if result.trim().is_empty() {
2263 summary.to_string()
2264 } else {
2265 result
2266 }
2267}
2268
2269fn cap_sentences(text: &str, n: usize) -> &str {
2273 let bytes = text.as_bytes();
2274 let len = bytes.len();
2275 let mut count = 0;
2276 let mut i = 0;
2277 while i < len {
2278 if bytes[i] == b'.' {
2280 let next = i + 1;
2281 if next < len && (bytes[next] == b' ' || bytes[next] == b'\n') {
2282 count += 1;
2283 if count >= n {
2284 return text[..=i].trim_end();
2286 }
2287 }
2288 }
2289 i += 1;
2290 }
2291 text.trim_end()
2293}
2294
2295pub(crate) fn excerpt_pub(text: &str) -> String {
2298 excerpt(text)
2299}
2300
2301fn excerpt(text: &str) -> String {
2302 let value = one_line(text);
2303 value.chars().take(256).collect()
2304}
2305
2306fn one_line(text: &str) -> String {
2307 text.split_whitespace().collect::<Vec<_>>().join(" ")
2308}
2309
2310const FILE_EXPAND_CAP_BYTES: usize = 2048;
2317
2318pub fn resolve_capsule(
2330 conn: &Connection,
2331 repo_root: &std::path::Path,
2332 handle: &str,
2333) -> kimetsu_core::KimetsuResult<String> {
2334 if let Some(memory_id) = handle.strip_prefix("memory:") {
2335 let mut stmt = conn.prepare_cached(
2337 "SELECT text FROM memories WHERE memory_id = ? AND invalidated_at IS NULL",
2338 )?;
2339 let text: Option<String> = stmt
2340 .query_row(rusqlite::params![memory_id], |row| row.get(0))
2341 .optional()?;
2342 match text {
2343 Some(t) => Ok(t),
2344 None => {
2345 Err(format!("expand_capsule: no active memory found for handle `{handle}`").into())
2346 }
2347 }
2348 } else if let Some(rel_path) = handle.strip_prefix("file:") {
2349 let path = std::path::Path::new(rel_path);
2354 if path.is_absolute() {
2355 return Err(format!(
2356 "expand_capsule: `{handle}` is an absolute path — only repo-relative paths are supported"
2357 )
2358 .into());
2359 }
2360 for component in path.components() {
2361 match component {
2362 std::path::Component::ParentDir => {
2363 return Err(format!(
2364 "expand_capsule: `{handle}` contains `..` traversal — rejected"
2365 )
2366 .into());
2367 }
2368 std::path::Component::RootDir | std::path::Component::Prefix(_) => {
2369 return Err(format!(
2370 "expand_capsule: `{handle}` is an absolute path — only repo-relative paths are supported"
2371 )
2372 .into());
2373 }
2374 _ => {}
2375 }
2376 }
2377 let full_path = repo_root.join(path);
2378 let bytes = std::fs::read(&full_path)
2379 .map_err(|e| format!("expand_capsule: could not read `{rel_path}`: {e}"))?;
2380 let bounded = if bytes.len() > FILE_EXPAND_CAP_BYTES {
2382 let mut end = FILE_EXPAND_CAP_BYTES;
2383 while end > 0 && (bytes[end] & 0xC0) == 0x80 {
2385 end -= 1;
2386 }
2387 let s = String::from_utf8_lossy(&bytes[..end]);
2388 format!(
2389 "{s}\n[... truncated at {FILE_EXPAND_CAP_BYTES} bytes; call expand_capsule again with a line range if needed]"
2390 )
2391 } else {
2392 String::from_utf8_lossy(&bytes).into_owned()
2393 };
2394 Ok(bounded)
2395 } else if handle.starts_with("run:") {
2396 Err(format!(
2397 "expand_capsule: `run:` handle expansion is not yet supported (handle: `{handle}`)"
2398 )
2399 .into())
2400 } else {
2401 Err(format!(
2402 "expand_capsule: unrecognised handle format `{handle}`; \
2403 expected `memory:<id>`, `file:<path>`, or `run:<id>`"
2404 )
2405 .into())
2406 }
2407}
2408
2409pub fn rerank_capsules(
2418 query: &str,
2419 capsules: Vec<ContextCapsule>,
2420 reranker: &dyn crate::embeddings::Reranker,
2421 floor: f32,
2422 cap: usize,
2423) -> Vec<ContextCapsule> {
2424 if capsules.is_empty() {
2425 return capsules;
2426 }
2427
2428 let docs: Vec<&str> = capsules.iter().map(|c| c.summary.as_str()).collect();
2434 let scores = match reranker.rerank(query, &docs) {
2435 Ok(s) if s.len() == docs.len() => s,
2440 _ => {
2441 let mut out = capsules;
2443 if cap > 0 && out.len() > cap {
2444 out.truncate(cap);
2445 }
2446 return out;
2447 }
2448 };
2449
2450 let mut ranked: Vec<ContextCapsule> = capsules
2451 .into_iter()
2452 .zip(scores)
2453 .map(|(mut c, s)| {
2454 c.score = s;
2455 c
2456 })
2457 .collect();
2458
2459 ranked.sort_by(|a, b| {
2460 b.score
2461 .partial_cmp(&a.score)
2462 .unwrap_or(std::cmp::Ordering::Equal)
2463 });
2464
2465 ranked.retain(|c| c.score >= floor);
2466
2467 if cap > 0 && ranked.len() > cap {
2468 ranked.truncate(cap);
2469 }
2470
2471 ranked
2472}
2473
2474#[cfg(test)]
2475mod tests {
2476 use super::*;
2477
2478 fn capsule(kind: &str, summary: &str) -> ContextCapsule {
2479 ContextCapsule {
2480 id: "c".into(),
2481 kind: kind.into(),
2482 summary: summary.into(),
2483 token_estimate: 1,
2484 expansion_handle: "memory:x".into(),
2485 provenance: vec![],
2486 confidence: 1.0,
2487 freshness: 1.0,
2488 relevance: 1.0,
2489 scope_weight: 1.0,
2490 score: 1.0,
2491 }
2492 }
2493
2494 fn make_test_dir(tag: &str) -> std::path::PathBuf {
2497 use std::time::{SystemTime, UNIX_EPOCH};
2498 let ts = SystemTime::now()
2499 .duration_since(UNIX_EPOCH)
2500 .map(|d| d.subsec_nanos())
2501 .unwrap_or(0);
2502 let dir = std::env::temp_dir().join(format!("kbrain_test_{tag}_{ts}"));
2503 std::fs::create_dir_all(&dir).expect("create test dir");
2504 dir
2505 }
2506
2507 #[test]
2508 fn capsule_matches_kind_reads_memory_summary_prefix() {
2509 let mem = capsule("memory", "project:failure_pattern - linker not found");
2511 assert!(capsule_matches_kind(&mem, "failure_pattern"));
2512 assert!(!capsule_matches_kind(&mem, "command"));
2513 let repo = capsule("repo_file", "src/lib.rs:command - run build");
2515 assert!(capsule_matches_kind(&repo, "repo_file"));
2516 assert!(!capsule_matches_kind(&repo, "command"));
2517 }
2518
2519 #[test]
2522 fn usefulness_multiplier_neutral_at_zero_uses() {
2523 assert!((usefulness_multiplier(0.0, 0) - 1.0).abs() < f32::EPSILON);
2525 assert!((usefulness_multiplier(5.0, 0) - 1.0).abs() < f32::EPSILON);
2526 assert!((usefulness_multiplier(-5.0, 0) - 1.0).abs() < f32::EPSILON);
2527 }
2528
2529 #[test]
2533 fn usefulness_multiplier_blends_smoothly_in_transition() {
2534 let one_use = usefulness_multiplier(1.0, 1);
2537 assert!((one_use - 1.166_666_6).abs() < 1e-4, "got {one_use}");
2538 let two_uses = usefulness_multiplier(2.0, 2);
2541 assert!((two_uses - 1.333_333_4).abs() < 1e-4, "got {two_uses}");
2542 let two_uses_bad = usefulness_multiplier(-2.0, 2);
2544 assert!(
2546 (two_uses_bad - 0.666_666_7).abs() < 1e-4,
2547 "got {two_uses_bad}"
2548 );
2549 }
2550
2551 #[test]
2555 fn usefulness_multiplier_maps_ratio_onto_envelope() {
2556 assert!((usefulness_multiplier(5.0, 5) - 1.5).abs() < f32::EPSILON);
2558 assert!((usefulness_multiplier(-5.0, 5) - 0.5).abs() < f32::EPSILON);
2560 let mid = usefulness_multiplier(0.0, 6);
2562 assert!((mid - 1.0).abs() < f32::EPSILON, "got {mid}");
2563 let high = usefulness_multiplier(2.0, 4);
2565 assert!((high - 1.25).abs() < f32::EPSILON, "got {high}");
2566 let low = usefulness_multiplier(-2.0, 4);
2568 assert!((low - 0.75).abs() < f32::EPSILON, "got {low}");
2569 }
2570
2571 #[test]
2575 fn usefulness_multiplier_clamps_to_envelope() {
2576 assert!((usefulness_multiplier(100.0, 5) - 1.5).abs() < f32::EPSILON);
2578 assert!((usefulness_multiplier(-100.0, 5) - 0.5).abs() < f32::EPSILON);
2580 }
2581
2582 #[test]
2585 fn query_tokens_expands_build_class() {
2586 let toks = query_tokens("Build the project from source");
2587 assert!(toks.iter().any(|t| t == "build"));
2588 assert!(toks.iter().any(|t| t == "shell_background"));
2590 assert!(toks.iter().any(|t| t == "long_running"));
2591 }
2592
2593 #[test]
2594 fn query_tokens_expands_edit_class() {
2595 let toks = query_tokens("Modify the config to fix the bug");
2596 assert!(toks.iter().any(|t| t == "edit_file"));
2597 assert!(toks.iter().any(|t| t == "apply_patch"));
2598 }
2599
2600 #[test]
2601 fn query_tokens_expands_search_class() {
2602 let toks = query_tokens("Find all references to the symbol");
2603 assert!(toks.iter().any(|t| t == "glob"));
2604 assert!(toks.iter().any(|t| t == "search_files"));
2605 }
2606
2607 #[test]
2608 fn query_tokens_no_expansion_on_unrelated_query() {
2609 let toks = query_tokens("hello world testing nothing");
2610 assert!(toks.iter().any(|t| t == "hello"));
2612 assert!(toks.iter().any(|t| t == "world"));
2614 }
2615
2616 #[test]
2619 fn jaccard_is_zero_for_disjoint_sets() {
2620 let a: std::collections::HashSet<String> =
2621 ["foo", "bar"].iter().map(|s| s.to_string()).collect();
2622 let b: std::collections::HashSet<String> =
2623 ["baz", "qux"].iter().map(|s| s.to_string()).collect();
2624 assert!((jaccard(&a, &b) - 0.0).abs() < f32::EPSILON);
2625 }
2626
2627 #[test]
2628 fn jaccard_is_one_for_identical_sets() {
2629 let a: std::collections::HashSet<String> =
2630 ["foo", "bar"].iter().map(|s| s.to_string()).collect();
2631 let b = a.clone();
2632 assert!((jaccard(&a, &b) - 1.0).abs() < f32::EPSILON);
2633 }
2634
2635 #[test]
2636 fn jaccard_partial_overlap() {
2637 let a: std::collections::HashSet<String> = ["foo", "bar", "baz"]
2638 .iter()
2639 .map(|s| s.to_string())
2640 .collect();
2641 let b: std::collections::HashSet<String> =
2642 ["bar", "qux"].iter().map(|s| s.to_string()).collect();
2643 assert!((jaccard(&a, &b) - 0.25).abs() < f32::EPSILON);
2645 }
2646
2647 #[test]
2648 fn summary_token_set_lowercases_and_filters_short() {
2649 let set = summary_token_set("Build the Foo-bar project");
2650 assert!(set.contains("build"));
2651 assert!(set.contains("foo"));
2652 assert!(set.contains("bar"));
2653 assert!(set.contains("project"));
2654 assert!(set.contains("the"));
2656 }
2657
2658 fn insert_memory_with_embedding(
2664 conn: &rusqlite::Connection,
2665 memory_id: &str,
2666 text: &str,
2667 embedder: &dyn embeddings::Embedder,
2668 ) {
2669 let normalized = kimetsu_core::memory::normalize_memory_text(text);
2670 conn.execute(
2671 "
2672 INSERT INTO memories (
2673 memory_id, scope, kind, text, normalized_text, confidence,
2674 source_event_id, provenance_snapshot_json, created_at,
2675 use_count, usefulness_score, embedding, embedding_model
2676 )
2677 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
2678 '2026-05-01T00:00:00Z', 0, 0.0, ?4, ?5)
2679 ",
2680 rusqlite::params![
2681 memory_id,
2682 text,
2683 normalized,
2684 embeddings::encode_embedding(&embedder.embed(text).expect("embed test row")),
2685 embedder.model_id(),
2686 ],
2687 )
2688 .expect("insert memory");
2689 conn.execute(
2690 "INSERT INTO memories_fts (memory_id, text, kind, scope) VALUES (?1, ?2, 'fact', 'global_user')",
2691 rusqlite::params![memory_id, text],
2692 )
2693 .expect("insert fts row");
2694 }
2695
2696 #[test]
2706 fn hybrid_retrieval_uses_cosine_score_to_rerank() {
2707 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
2708 crate::schema::initialize(&conn).expect("init schema");
2709 let stub = embeddings::StubEmbedder::new();
2710
2711 insert_memory_with_embedding(&conn, "m_rg", "use ripgrep for code search", &stub);
2712 insert_memory_with_embedding(
2713 &conn,
2714 "m_unrelated",
2715 "cookie recipe with chocolate chips",
2716 &stub,
2717 );
2718
2719 let weights = kimetsu_core::config::BrokerWeights::default();
2722 let bundle = retrieve_context_with_embedder(
2723 &conn,
2724 "/fake-repo",
2725 &weights,
2726 ContextRequest {
2727 stage: "localization".to_string(),
2728 query: "ripgrep search".to_string(),
2729 budget_tokens: 4000,
2730 ..Default::default()
2731 },
2732 &[],
2733 &stub,
2734 )
2735 .expect("retrieve");
2736
2737 let memory_handles: Vec<_> = bundle
2738 .capsules
2739 .iter()
2740 .filter(|c| c.expansion_handle.starts_with("memory:"))
2741 .collect();
2742 assert!(
2743 !memory_handles.is_empty(),
2744 "at least one memory should surface"
2745 );
2746 assert_eq!(
2748 memory_handles[0].expansion_handle,
2749 "memory:m_rg",
2750 "ripgrep memory should outrank the cookie recipe; ranked: {:?}",
2751 memory_handles
2752 .iter()
2753 .map(|c| &c.expansion_handle)
2754 .collect::<Vec<_>>()
2755 );
2756 }
2757
2758 #[test]
2765 fn hybrid_retrieval_skips_cosine_on_model_id_mismatch() {
2766 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
2767 crate::schema::initialize(&conn).expect("init schema");
2768 let stub = embeddings::StubEmbedder::new();
2769 insert_memory_with_embedding(&conn, "m_xref", "use ripgrep for code search", &stub);
2770
2771 conn.execute(
2776 "UPDATE memories SET embedding_model = 'bge-small-en-v1.5' WHERE memory_id = 'm_xref'",
2777 [],
2778 )
2779 .expect("force model_id mismatch");
2780
2781 let weights = kimetsu_core::config::BrokerWeights::default();
2786 let bundle = retrieve_context_with_embedder(
2787 &conn,
2788 "/fake-repo",
2789 &weights,
2790 ContextRequest {
2791 stage: "localization".to_string(),
2792 query: "ripgrep search".to_string(),
2793 budget_tokens: 4000,
2794 ..Default::default()
2795 },
2796 &[],
2797 &stub,
2798 )
2799 .expect("retrieve");
2800
2801 assert!(
2802 bundle
2803 .capsules
2804 .iter()
2805 .any(|c| c.expansion_handle == "memory:m_xref"),
2806 "cross-model row should still match lexically (cosine skipped, FTS works)"
2807 );
2808 }
2809
2810 #[test]
2817 fn usefulness_decay_disabled_when_half_life_is_zero_or_negative() {
2818 let ancient = "2021-01-01T00:00:00Z";
2820 assert!((usefulness_decay(Some(ancient), ancient, 0.0) - 1.0).abs() < f32::EPSILON);
2821 assert!((usefulness_decay(Some(ancient), ancient, -1.0) - 1.0).abs() < f32::EPSILON);
2822 }
2823
2824 #[test]
2828 fn usefulness_decay_returns_one_on_unparseable_timestamps() {
2829 assert!(
2830 (usefulness_decay(Some("not-a-date"), "also-not", 30.0) - 1.0).abs() < f32::EPSILON
2831 );
2832 }
2833
2834 #[test]
2837 fn usefulness_decay_full_at_zero_age() {
2838 let future = "2099-01-01T00:00:00Z";
2840 let d = usefulness_decay(Some(future), future, 30.0);
2841 assert!((d - 1.0).abs() < f32::EPSILON, "got {d}");
2842 }
2843
2844 #[test]
2849 fn usefulness_decay_follows_half_life_curve() {
2850 let half_life = 10.0_f32;
2851 let now = OffsetDateTime::now_utc();
2852 let fmt = &time::format_description::well_known::Rfc3339;
2853
2854 let one_half_life_ago = (now - time::Duration::seconds((half_life * 86_400.0) as i64))
2856 .format(fmt)
2857 .expect("format");
2858 let d1 = usefulness_decay(Some(&one_half_life_ago), &one_half_life_ago, half_life);
2859 assert!(
2860 (d1 - 0.5).abs() < 0.01,
2861 "expected ~0.5 at one half-life, got {d1}"
2862 );
2863
2864 let two_half_lives_ago = (now
2866 - time::Duration::seconds((2.0 * half_life * 86_400.0) as i64))
2867 .format(fmt)
2868 .expect("format");
2869 let d2 = usefulness_decay(Some(&two_half_lives_ago), &two_half_lives_ago, half_life);
2870 assert!(
2871 (d2 - 0.25).abs() < 0.01,
2872 "expected ~0.25 at two half-lives, got {d2}"
2873 );
2874 }
2875
2876 #[test]
2880 fn usefulness_decay_falls_back_to_created_at_when_last_useful_is_none() {
2881 let now = OffsetDateTime::now_utc();
2882 let fmt = &time::format_description::well_known::Rfc3339;
2883 let one_day_ago = (now - time::Duration::seconds(86_400))
2884 .format(fmt)
2885 .expect("format");
2886 let d = usefulness_decay(None, &one_day_ago, 30.0);
2887 assert!(
2889 (d - 0.977).abs() < 0.01,
2890 "expected ~0.977 for 1-day-old created_at under 30d half-life, got {d}"
2891 );
2892 }
2893
2894 #[test]
2899 fn aged_cited_memory_ranks_below_recently_cited_memory() {
2900 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
2901 crate::schema::initialize(&conn).expect("init schema");
2902
2903 let now = OffsetDateTime::now_utc();
2904 let fmt = &time::format_description::well_known::Rfc3339;
2905 let one_day_ago = (now - time::Duration::seconds(86_400))
2906 .format(fmt)
2907 .expect("format");
2908 let one_year_ago = (now - time::Duration::seconds(365 * 86_400))
2909 .format(fmt)
2910 .expect("format");
2911
2912 for (mid, last_useful) in [("m_recent", &one_day_ago), ("m_aged", &one_year_ago)] {
2916 let text = "use ripgrep for code search";
2917 let normalized = kimetsu_core::memory::normalize_memory_text(text);
2918 conn.execute(
2919 "
2920 INSERT INTO memories (
2921 memory_id, scope, kind, text, normalized_text, confidence,
2922 source_event_id, provenance_snapshot_json, created_at,
2923 use_count, usefulness_score, last_useful_at
2924 )
2925 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
2926 '2024-01-01T00:00:00Z', 5, 5.0, ?4)
2927 ",
2928 rusqlite::params![mid, text, normalized, last_useful],
2929 )
2930 .expect("insert memory");
2931 conn.execute(
2932 "INSERT INTO memories_fts (memory_id, text, kind, scope)
2933 VALUES (?1, ?2, 'fact', 'global_user')",
2934 rusqlite::params![mid, text],
2935 )
2936 .expect("insert fts");
2937 }
2938
2939 let weights = kimetsu_core::config::BrokerWeights::default();
2941 let bundle = retrieve_context_with_embedder(
2942 &conn,
2943 "/fake-repo",
2944 &weights,
2945 ContextRequest {
2946 stage: "localization".to_string(),
2947 query: "ripgrep search".to_string(),
2948 budget_tokens: 4000,
2949 ..Default::default()
2950 },
2951 &[],
2952 &embeddings::NoopEmbedder,
2953 )
2954 .expect("retrieve");
2955
2956 let mem_order: Vec<&str> = bundle
2957 .capsules
2958 .iter()
2959 .filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
2960 .collect();
2961 assert_eq!(
2962 mem_order.first().copied(),
2963 Some("m_recent"),
2964 "recently-cited memory must rank first under decay; got order {mem_order:?}"
2965 );
2966 }
2967
2968 #[test]
2973 fn aged_cited_memory_does_not_decay_when_half_life_is_zero() {
2974 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
2975 crate::schema::initialize(&conn).expect("init schema");
2976
2977 let now = OffsetDateTime::now_utc();
2978 let fmt = &time::format_description::well_known::Rfc3339;
2979 let one_day_ago = (now - time::Duration::seconds(86_400))
2980 .format(fmt)
2981 .expect("format");
2982 let one_year_ago = (now - time::Duration::seconds(365 * 86_400))
2983 .format(fmt)
2984 .expect("format");
2985
2986 for (mid, last_useful) in [("m_recent", &one_day_ago), ("m_aged", &one_year_ago)] {
2987 let text = "use ripgrep for code search";
2988 let normalized = kimetsu_core::memory::normalize_memory_text(text);
2989 conn.execute(
2990 "
2991 INSERT INTO memories (
2992 memory_id, scope, kind, text, normalized_text, confidence,
2993 source_event_id, provenance_snapshot_json, created_at,
2994 use_count, usefulness_score, last_useful_at
2995 )
2996 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
2997 '2024-01-01T00:00:00Z', 5, 5.0, ?4)
2998 ",
2999 rusqlite::params![mid, text, normalized, last_useful],
3000 )
3001 .expect("insert memory");
3002 conn.execute(
3003 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3004 VALUES (?1, ?2, 'fact', 'global_user')",
3005 rusqlite::params![mid, text],
3006 )
3007 .expect("insert fts");
3008 }
3009
3010 let weights = kimetsu_core::config::BrokerWeights {
3012 decay_half_life_days: 0.0,
3013 ..Default::default()
3014 };
3015
3016 let bundle = retrieve_context_with_embedder(
3017 &conn,
3018 "/fake-repo",
3019 &weights,
3020 ContextRequest {
3021 stage: "localization".to_string(),
3022 query: "ripgrep search".to_string(),
3023 budget_tokens: 4000,
3024 ..Default::default()
3025 },
3026 &[],
3027 &embeddings::NoopEmbedder,
3028 )
3029 .expect("retrieve");
3030
3031 let scores: Vec<(String, f32)> = bundle
3037 .capsules
3038 .iter()
3039 .filter_map(|c| {
3040 c.expansion_handle
3041 .strip_prefix("memory:")
3042 .map(|id| (id.to_string(), c.score))
3043 })
3044 .collect();
3045 assert_eq!(scores.len(), 2, "both memories should surface");
3046 let recent_score = scores
3047 .iter()
3048 .find(|(id, _)| id == "m_recent")
3049 .map(|(_, s)| *s)
3050 .expect("m_recent present");
3051 let aged_score = scores
3052 .iter()
3053 .find(|(id, _)| id == "m_aged")
3054 .map(|(_, s)| *s)
3055 .expect("m_aged present");
3056 assert!(
3058 (recent_score - aged_score).abs() < 1e-4,
3059 "with decay disabled the two memories should tie on score: recent={recent_score} aged={aged_score}"
3060 );
3061 }
3062
3063 #[test]
3068 fn hybrid_retrieval_with_noop_embedder_is_lexical_only() {
3069 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3070 crate::schema::initialize(&conn).expect("init schema");
3071 let stub = embeddings::StubEmbedder::new();
3072 insert_memory_with_embedding(&conn, "m_a", "use ripgrep", &stub);
3074 insert_memory_with_embedding(&conn, "m_b", "use ripgrep too", &stub);
3075
3076 let weights = kimetsu_core::config::BrokerWeights::default();
3079 let bundle = retrieve_context_with_embedder(
3080 &conn,
3081 "/fake-repo",
3082 &weights,
3083 ContextRequest {
3084 stage: "localization".to_string(),
3085 query: "ripgrep".to_string(),
3086 budget_tokens: 4000,
3087 ..Default::default()
3088 },
3089 &[],
3090 &embeddings::NoopEmbedder,
3091 )
3092 .expect("retrieve");
3093
3094 let count = bundle
3095 .capsules
3096 .iter()
3097 .filter(|c| c.expansion_handle.starts_with("memory:"))
3098 .count();
3099 assert_eq!(count, 2, "both memories should surface via FTS");
3100 }
3101
3102 #[cfg(feature = "embeddings")]
3129 #[test]
3130 fn ann_finds_semantic_match_fts_misses() {
3131 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3132 crate::schema::initialize(&conn).expect("init schema");
3133
3134 struct OracleEmbedder;
3137 impl embeddings::Embedder for OracleEmbedder {
3138 fn embed(&self, _text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
3139 Ok(vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
3141 }
3142 fn model_id(&self) -> &str {
3143 "oracle-d8"
3144 }
3145 fn dim(&self) -> usize {
3146 8
3147 }
3148 }
3149
3150 let model_id = "oracle-d8";
3151
3152 let sem_vec = embeddings::encode_embedding(&[1.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
3155 let sem_text = "cookie recipe chocolate";
3156 let sem_norm = kimetsu_core::memory::normalize_memory_text(sem_text);
3157 conn.execute(
3158 "INSERT INTO memories (
3159 memory_id, scope, kind, text, normalized_text, confidence,
3160 source_event_id, provenance_snapshot_json, created_at,
3161 use_count, usefulness_score, embedding, embedding_model
3162 )
3163 VALUES ('m_semantic', 'global_user', 'fact', ?1, ?2, 1.0, NULL, '{}',
3164 '2026-01-01T00:00:00Z', 0, 0.0, ?3, ?4)",
3165 rusqlite::params![sem_text, sem_norm, sem_vec, model_id],
3166 )
3167 .expect("insert m_semantic");
3168 conn.execute(
3169 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3170 VALUES ('m_semantic', ?1, 'fact', 'global_user')",
3171 rusqlite::params![sem_text],
3172 )
3173 .expect("insert m_semantic fts");
3174
3175 let decoy_vec = embeddings::encode_embedding(&[0.0f32, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
3177 let decoy_text = "git rebase squash commits";
3178 let decoy_norm = kimetsu_core::memory::normalize_memory_text(decoy_text);
3179 conn.execute(
3180 "INSERT INTO memories (
3181 memory_id, scope, kind, text, normalized_text, confidence,
3182 source_event_id, provenance_snapshot_json, created_at,
3183 use_count, usefulness_score, embedding, embedding_model
3184 )
3185 VALUES ('m_decoy', 'global_user', 'fact', ?1, ?2, 1.0, NULL, '{}',
3186 '2026-01-01T00:00:00Z', 0, 0.0, ?3, ?4)",
3187 rusqlite::params![decoy_text, decoy_norm, decoy_vec, model_id],
3188 )
3189 .expect("insert m_decoy");
3190 conn.execute(
3191 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3192 VALUES ('m_decoy', ?1, 'fact', 'global_user')",
3193 rusqlite::params![decoy_text],
3194 )
3195 .expect("insert m_decoy fts");
3196
3197 let fts_hits: i64 = conn
3199 .query_row(
3200 "SELECT COUNT(*) FROM memories_fts \
3201 WHERE memories_fts MATCH 'phosphorescent bioluminescent'",
3202 [],
3203 |r| r.get(0),
3204 )
3205 .unwrap_or(0);
3206 assert_eq!(
3207 fts_hits, 0,
3208 "sanity: query tokens must not appear in any memory text"
3209 );
3210
3211 let weights = kimetsu_core::config::BrokerWeights::default();
3215 let bundle = retrieve_context_with_embedder(
3216 &conn,
3217 "/fake-repo",
3218 &weights,
3219 ContextRequest {
3220 stage: "localization".to_string(),
3221 query: "phosphorescent bioluminescent organism".to_string(),
3222 budget_tokens: 4000,
3223 ..Default::default()
3224 },
3225 &[],
3226 &OracleEmbedder,
3227 )
3228 .expect("retrieve");
3229
3230 let handles: Vec<&str> = bundle
3231 .capsules
3232 .iter()
3233 .filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
3234 .collect();
3235
3236 assert!(
3237 handles.contains(&"m_semantic"),
3238 "ANN must surface m_semantic (cosine=1 with oracle query) even though \
3239 FTS found nothing; got handles: {handles:?}"
3240 );
3241 }
3242
3243 #[cfg(feature = "embeddings")]
3245 #[test]
3246 fn dedup_memory_matched_by_fts_and_ann_appears_once() {
3247 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3248 crate::schema::initialize(&conn).expect("init schema");
3249
3250 let stub = embeddings::StubEmbedder::new();
3251
3252 insert_memory_with_embedding(&conn, "m_both", "use ripgrep for fast search", &stub);
3256
3257 let weights = kimetsu_core::config::BrokerWeights::default();
3258 let bundle = retrieve_context_with_embedder(
3259 &conn,
3260 "/fake-repo",
3261 &weights,
3262 ContextRequest {
3263 stage: "localization".to_string(),
3264 query: "ripgrep".to_string(),
3265 budget_tokens: 4000,
3266 ..Default::default()
3267 },
3268 &[],
3269 &stub,
3270 )
3271 .expect("retrieve");
3272
3273 let count = bundle
3274 .capsules
3275 .iter()
3276 .filter(|c| c.expansion_handle == "memory:m_both")
3277 .count();
3278 assert_eq!(
3279 count,
3280 1,
3281 "m_both (matched by both FTS and ANN) must appear exactly once; \
3282 bundle: {:?}",
3283 bundle
3284 .capsules
3285 .iter()
3286 .map(|c| &c.expansion_handle)
3287 .collect::<Vec<_>>()
3288 );
3289 }
3290
3291 #[cfg(feature = "embeddings")]
3315 #[test]
3316 fn embedding_mmr_collapses_paraphrases_but_jaccard_does_not() {
3317 struct OracleEmbedder;
3320 impl embeddings::Embedder for OracleEmbedder {
3321 fn embed(&self, _text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
3322 let mut v = vec![0.0f32; 8];
3323 v[0] = 1.0;
3324 Ok(v)
3325 }
3326 fn model_id(&self) -> &str {
3327 "oracle-d8"
3328 }
3329 fn dim(&self) -> usize {
3330 8
3331 }
3332 }
3333
3334 let oracle = OracleEmbedder;
3337 let weights = kimetsu_core::config::BrokerWeights::default();
3338
3339 let m_rg1_text = "prefer ripgrep for searching source code";
3342 let m_rg2_text = "rg is the fastest way to locate patterns";
3343
3344 let conn = rusqlite::Connection::open_in_memory().expect("in-memory");
3350 crate::schema::initialize(&conn).expect("init schema");
3351 insert_memory_with_embedding(&conn, "m_rg1", m_rg1_text, &oracle);
3352 insert_memory_with_embedding(&conn, "m_rg2", m_rg2_text, &oracle);
3353
3354 let bundle_embedding = retrieve_context_with_embedder(
3355 &conn,
3356 "/fake-repo",
3357 &weights,
3358 ContextRequest {
3359 stage: "localization".to_string(),
3360 query: "search source patterns".to_string(),
3362 budget_tokens: 20_000,
3363 max_capsules: 1, ..Default::default()
3365 },
3366 &[],
3367 &oracle,
3368 )
3369 .expect("retrieve with oracle embedder");
3370
3371 let emb_in_capsules = bundle_embedding
3374 .capsules
3375 .iter()
3376 .filter(|c| {
3377 c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
3378 })
3379 .count();
3380 assert_eq!(
3381 emb_in_capsules,
3382 1,
3383 "embedding-MMR must collapse cosine=1.0 paraphrases: with max_capsules=1 \
3384 only ONE should be included; capsule handles: {:?}; excluded: {:?}",
3385 bundle_embedding
3386 .capsules
3387 .iter()
3388 .map(|c| &c.expansion_handle)
3389 .collect::<Vec<_>>(),
3390 bundle_embedding
3391 .excluded
3392 .iter()
3393 .map(|c| &c.expansion_handle)
3394 .collect::<Vec<_>>()
3395 );
3396
3397 let emb_in_excluded = bundle_embedding
3399 .excluded
3400 .iter()
3401 .filter(|c| {
3402 c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
3403 })
3404 .count();
3405 assert_eq!(
3406 emb_in_excluded,
3407 1,
3408 "the second near-duplicate must be in excluded under embedding-MMR; \
3409 excluded handles: {:?}",
3410 bundle_embedding
3411 .excluded
3412 .iter()
3413 .map(|c| &c.expansion_handle)
3414 .collect::<Vec<_>>()
3415 );
3416
3417 let conn2 = rusqlite::Connection::open_in_memory().expect("in-memory 2");
3422 crate::schema::initialize(&conn2).expect("init schema 2");
3423 insert_memory_with_embedding(&conn2, "m_rg1", m_rg1_text, &oracle);
3424 insert_memory_with_embedding(&conn2, "m_rg2", m_rg2_text, &oracle);
3425
3426 let bundle_lean = retrieve_context_with_embedder(
3427 &conn2,
3428 "/fake-repo",
3429 &weights,
3430 ContextRequest {
3431 stage: "localization".to_string(),
3432 query: "search source patterns".to_string(),
3433 budget_tokens: 20_000,
3434 max_capsules: 2, ..Default::default()
3436 },
3437 &[],
3438 &embeddings::NoopEmbedder,
3439 )
3440 .expect("retrieve with NoopEmbedder");
3441
3442 let lean_in_capsules = bundle_lean
3443 .capsules
3444 .iter()
3445 .filter(|c| {
3446 c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
3447 })
3448 .count();
3449 assert_eq!(
3450 lean_in_capsules,
3451 2,
3452 "Jaccard-only path must NOT collapse the two paraphrases (different words, \
3453 low token overlap → both survive MMR with max_capsules=2); capsule handles: {:?}",
3454 bundle_lean
3455 .capsules
3456 .iter()
3457 .map(|c| &c.expansion_handle)
3458 .collect::<Vec<_>>()
3459 );
3460 }
3461
3462 #[test]
3465 fn content_tokens_strips_stopwords_keeps_topical_words() {
3466 let got = content_tokens("Tell me about kimetsu, what's the idea of the repo");
3467 assert_eq!(got, vec!["kimetsu", "idea", "repo"]);
3470 }
3471
3472 #[test]
3473 fn light_stem_strips_one_inflection_suffix() {
3474 assert_eq!(light_stem("benchmarked"), "benchmark");
3475 assert_eq!(light_stem("benchmarking"), "benchmark");
3476 assert_eq!(light_stem("repos"), "repo");
3477 assert_eq!(light_stem("does"), "does");
3479 assert_eq!(light_stem("toml"), "toml");
3480 }
3481
3482 #[test]
3489 fn stemmed_query_matches_inflected_corpus_through_floor() {
3490 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3491 crate::schema::initialize(&conn).expect("init schema");
3492 let insert = |id: &str, text: &str| {
3493 let norm = kimetsu_core::memory::normalize_memory_text(text);
3494 conn.execute(
3495 "INSERT INTO memories (
3496 memory_id, scope, kind, text, normalized_text, confidence,
3497 source_event_id, provenance_snapshot_json, created_at,
3498 use_count, usefulness_score, embedding, embedding_model
3499 )
3500 VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
3501 '2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
3502 rusqlite::params![id, text, norm],
3503 )
3504 .expect("insert memory");
3505 conn.execute(
3506 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3507 VALUES (?1, ?2, 'fact', 'global_user')",
3508 rusqlite::params![id, text],
3509 )
3510 .expect("insert fts");
3511 };
3512 insert(
3513 "m_bench",
3514 "kimetsu benchmark runs go through the kbench binary and the Terminal-Bench driver",
3515 );
3516 insert(
3517 "m_doctor",
3518 "kimetsu doctor version-skew check parses process start times on Windows via CIM",
3519 );
3520 insert(
3521 "m_gc",
3522 "kimetsu runs auto-GC on run creation; keep the env guard at the trigger site",
3523 );
3524
3525 let bundle = retrieve_context_with_embedder(
3526 &conn,
3527 "/fake-repo",
3528 &kimetsu_core::config::BrokerWeights::default(),
3529 ContextRequest {
3530 stage: "localization".to_string(),
3531 query: "Can you find out how kimetsu is benchmarked?".to_string(),
3532 budget_tokens: 2000,
3533 max_capsules: 2,
3534 min_lexical_coverage: 0.5,
3535 ..Default::default()
3536 },
3537 &[],
3538 &embeddings::NoopEmbedder,
3539 )
3540 .expect("retrieve");
3541 let handles: Vec<_> = bundle
3542 .capsules
3543 .iter()
3544 .map(|c| c.expansion_handle.as_str())
3545 .collect();
3546 assert!(
3547 handles.contains(&"memory:m_bench"),
3548 "stemmed 'benchmarked' must surface the benchmark memory; got {handles:?}"
3549 );
3550 assert!(
3551 !handles.contains(&"memory:m_doctor") && !handles.contains(&"memory:m_gc"),
3552 "off-topic memories sharing only 'kimetsu' must stay below the floor; got {handles:?}"
3553 );
3554 }
3555
3556 #[test]
3557 fn weighted_coverage_ignores_zero_idf_tokens() {
3558 let content = vec![
3562 "kimetsu".to_string(),
3563 "idea".to_string(),
3564 "repo".to_string(),
3565 ];
3566 let mut idf = HashMap::new();
3567 idf.insert("kimetsu".to_string(), 0.0);
3568 idf.insert("idea".to_string(), 1.386);
3569 idf.insert("repo".to_string(), 0.693);
3570
3571 let cov = weighted_coverage(
3573 &content,
3574 &idf,
3575 "global:fact - the git repo and kimetsu brain",
3576 );
3577 assert!((cov - 0.333).abs() < 0.01, "got {cov}");
3578
3579 let cov_topical =
3581 weighted_coverage(&content, &idf, "global:fact - the core idea of kimetsu");
3582 assert!(cov_topical > 0.6, "got {cov_topical}");
3583 }
3584
3585 #[test]
3586 fn escape_like_neutralizes_wildcards() {
3587 assert_eq!(escape_like("a_b%c"), "a\\_b\\%c");
3588 assert_eq!(escape_like("plain"), "plain");
3589 }
3590
3591 #[test]
3605 fn lexical_floor_drops_offtopic_memories_sharing_project_name() {
3606 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3607 crate::schema::initialize(&conn).expect("init schema");
3608
3609 let insert = |id: &str, text: &str| {
3610 let norm = kimetsu_core::memory::normalize_memory_text(text);
3611 conn.execute(
3612 "INSERT INTO memories (
3613 memory_id, scope, kind, text, normalized_text, confidence,
3614 source_event_id, provenance_snapshot_json, created_at,
3615 use_count, usefulness_score, embedding, embedding_model
3616 )
3617 VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
3618 '2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
3619 rusqlite::params![id, text, norm],
3620 )
3621 .expect("insert memory");
3622 conn.execute(
3623 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3624 VALUES (?1, ?2, 'fact', 'global_user')",
3625 rusqlite::params![id, text],
3626 )
3627 .expect("insert fts");
3628 };
3629
3630 insert(
3633 "m1",
3634 "When implementing a setup command that calls init_project, tests must call \
3635 git_init_boundary before setup_cmd so ProjectPaths discover resolves to the temp \
3636 dir instead of climbing to the real parent git repo including the user brain at kimetsu",
3637 );
3638 insert(
3639 "m2",
3640 "A member crate with default embeddings silently turned embeddings on for the entire \
3641 cargo test workspace build graph because cargo unifies features; kimetsu-chat \
3642 retrieval tests failed",
3643 );
3644 insert(
3645 "m3",
3646 "In toml 0.9 use toml from_str to parse a TOML document into a Value not str parse; \
3647 implementing config get and set in kimetsu-cli",
3648 );
3649
3650 let query = "Tell me about kimetsu, what's the idea of the repo".to_string();
3651 let weights = kimetsu_core::config::BrokerWeights::default();
3652 let handles = |bundle: &ContextBundle| {
3653 bundle
3654 .capsules
3655 .iter()
3656 .map(|c| c.expansion_handle.clone())
3657 .collect::<Vec<_>>()
3658 };
3659
3660 let no_floor = retrieve_context_with_embedder(
3662 &conn,
3663 "/fake-repo",
3664 &weights,
3665 ContextRequest {
3666 stage: "localization".to_string(),
3667 query: query.clone(),
3668 budget_tokens: 2000,
3669 max_capsules: 8,
3670 min_lexical_coverage: 0.0,
3671 ..Default::default()
3672 },
3673 &[],
3674 &embeddings::NoopEmbedder,
3675 )
3676 .expect("retrieve without floor");
3677 let before = handles(&no_floor);
3678 assert!(
3679 before.contains(&"memory:m2".to_string()) && before.contains(&"memory:m3".to_string()),
3680 "sanity: without the floor the pure-project-name memories should surface; got {before:?}"
3681 );
3682
3683 let floored = retrieve_context_with_embedder(
3685 &conn,
3686 "/fake-repo",
3687 &weights,
3688 ContextRequest {
3689 stage: "localization".to_string(),
3690 query,
3691 budget_tokens: 2000,
3692 max_capsules: 8,
3693 min_lexical_coverage: 0.5,
3694 ..Default::default()
3695 },
3696 &[],
3697 &embeddings::NoopEmbedder,
3698 )
3699 .expect("retrieve with floor");
3700 let after = handles(&floored);
3701 assert!(
3702 !after.contains(&"memory:m2".to_string()) && !after.contains(&"memory:m3".to_string()),
3703 "the lexical floor must drop memories whose only match is the corpus-ubiquitous \
3704 project name; surviving: {after:?}"
3705 );
3706 }
3707
3708 #[test]
3711 fn lexical_floor_keeps_ontopic_memory() {
3712 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3713 crate::schema::initialize(&conn).expect("init schema");
3714
3715 let insert = |id: &str, text: &str| {
3716 let norm = kimetsu_core::memory::normalize_memory_text(text);
3717 conn.execute(
3718 "INSERT INTO memories (
3719 memory_id, scope, kind, text, normalized_text, confidence,
3720 source_event_id, provenance_snapshot_json, created_at,
3721 use_count, usefulness_score, embedding, embedding_model
3722 )
3723 VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
3724 '2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
3725 rusqlite::params![id, text, norm],
3726 )
3727 .expect("insert memory");
3728 conn.execute(
3729 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3730 VALUES (?1, ?2, 'fact', 'global_user')",
3731 rusqlite::params![id, text],
3732 )
3733 .expect("insert fts");
3734 };
3735
3736 insert(
3738 "d1",
3739 "The distiller runs at session end and harvests durable lessons from the transcript",
3740 );
3741 insert(
3742 "n1",
3743 "Unrelated note about git rebase and squashing commits",
3744 );
3745
3746 let bundle = retrieve_context_with_embedder(
3747 &conn,
3748 "/fake-repo",
3749 &kimetsu_core::config::BrokerWeights::default(),
3750 ContextRequest {
3751 stage: "localization".to_string(),
3752 query: "how does the distiller work".to_string(),
3753 budget_tokens: 2000,
3754 min_lexical_coverage: 0.5,
3755 ..Default::default()
3756 },
3757 &[],
3758 &embeddings::NoopEmbedder,
3759 )
3760 .expect("retrieve");
3761
3762 assert!(
3763 bundle
3764 .capsules
3765 .iter()
3766 .any(|c| c.expansion_handle == "memory:d1"),
3767 "on-topic memory covering the rare query word must survive the floor; got: {:?}",
3768 bundle
3769 .capsules
3770 .iter()
3771 .map(|c| &c.expansion_handle)
3772 .collect::<Vec<_>>()
3773 );
3774 }
3775
3776 #[cfg(feature = "embeddings")]
3785 #[test]
3786 fn min_semantic_score_floor_drops_off_topic_queries() {
3787 struct DirectionalEmbedder {
3798 marker: &'static str,
3800 }
3801 impl embeddings::Embedder for DirectionalEmbedder {
3802 fn embed(&self, text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
3803 let mut v = vec![0.0f32; 8];
3804 if text.contains(self.marker) {
3805 v[0] = 1.0;
3806 } else {
3807 v[1] = 1.0;
3808 }
3809 Ok(v)
3810 }
3811 fn model_id(&self) -> &str {
3812 "directional-d8"
3813 }
3814 fn dim(&self) -> usize {
3815 8
3816 }
3817 }
3818
3819 let emb = DirectionalEmbedder { marker: "TOPIC_A" };
3820
3821 let conn = rusqlite::Connection::open_in_memory().expect("in-memory");
3822 crate::schema::initialize(&conn).expect("init schema");
3823
3824 insert_memory_with_embedding(&conn, "m_b", "cookie recipe chocolate baking TOPIC_B", &emb);
3826
3827 let weights = kimetsu_core::config::BrokerWeights::default();
3828
3829 let bundle_off = retrieve_context_with_embedder(
3831 &conn,
3832 "/fake-repo",
3833 &weights,
3834 ContextRequest {
3835 stage: "localization".to_string(),
3836 query: "TOPIC_A unrelated phosphorescent".to_string(),
3838 budget_tokens: 4000,
3839 min_semantic_score: 0.1, ..Default::default()
3841 },
3842 &[],
3843 &emb,
3844 )
3845 .expect("retrieve off-topic");
3846
3847 assert!(
3848 bundle_off.capsules.is_empty(),
3849 "off-topic query (cosine=0 < floor=0.1) must produce zero capsules; \
3850 got: {:?}",
3851 bundle_off
3852 .capsules
3853 .iter()
3854 .map(|c| &c.expansion_handle)
3855 .collect::<Vec<_>>()
3856 );
3857
3858 let conn2 = rusqlite::Connection::open_in_memory().expect("in-memory 2");
3861 crate::schema::initialize(&conn2).expect("init schema 2");
3862 insert_memory_with_embedding(
3863 &conn2,
3864 "m_b2",
3865 "cookie recipe chocolate TOPIC_B baking"
3866 .to_string()
3867 .as_str(),
3868 &emb,
3869 );
3870
3871 let bundle_on = retrieve_context_with_embedder(
3872 &conn2,
3873 "/fake-repo",
3874 &weights,
3875 ContextRequest {
3876 stage: "localization".to_string(),
3877 query: "cookie chocolate TOPIC_B".to_string(),
3879 budget_tokens: 4000,
3880 min_semantic_score: 0.1,
3881 ..Default::default()
3882 },
3883 &[],
3884 &emb,
3885 )
3886 .expect("retrieve on-topic");
3887
3888 assert!(
3889 bundle_on
3890 .capsules
3891 .iter()
3892 .any(|c| c.expansion_handle == "memory:m_b2"),
3893 "on-topic query (cosine=1.0 ≥ floor) must surface m_b2; \
3894 got capsules: {:?}",
3895 bundle_on
3896 .capsules
3897 .iter()
3898 .map(|c| &c.expansion_handle)
3899 .collect::<Vec<_>>()
3900 );
3901
3902 let conn3 = rusqlite::Connection::open_in_memory().expect("in-memory 3");
3905 crate::schema::initialize(&conn3).expect("init schema 3");
3906 insert_memory_with_embedding(
3907 &conn3,
3908 "m_b3",
3909 "cookie chocolate TOPIC_B recipe".to_string().as_str(),
3910 &emb,
3911 );
3912
3913 let bundle_noop_floor = retrieve_context_with_embedder(
3914 &conn3,
3915 "/fake-repo",
3916 &weights,
3917 ContextRequest {
3918 stage: "localization".to_string(),
3919 query: "cookie chocolate TOPIC_A".to_string(),
3921 budget_tokens: 4000,
3922 min_semantic_score: 0.0, ..Default::default()
3924 },
3925 &[],
3926 &emb,
3927 )
3928 .expect("retrieve noop floor");
3929
3930 assert!(
3932 bundle_noop_floor
3933 .capsules
3934 .iter()
3935 .any(|c| c.expansion_handle == "memory:m_b3"),
3936 "with floor=0.0 (disabled), off-topic-cosine memory must still surface via FTS; \
3937 got: {:?}",
3938 bundle_noop_floor
3939 .capsules
3940 .iter()
3941 .map(|c| &c.expansion_handle)
3942 .collect::<Vec<_>>()
3943 );
3944 }
3945
3946 #[cfg(feature = "embeddings")]
3975 #[test]
3976 fn d1f_token_economy_fewer_capsules_signal_preserved() {
3977 struct OracleTopicEmbedder;
3979 impl embeddings::Embedder for OracleTopicEmbedder {
3980 fn embed(&self, text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
3981 let mut v = vec![0.0f32; 8];
3982 if text.contains("TOPIC_A") {
3983 v[0] = 1.0; } else {
3985 v[1] = 1.0; }
3987 Ok(v)
3988 }
3989 fn model_id(&self) -> &str {
3990 "oracle-topic-d8"
3991 }
3992 fn dim(&self) -> usize {
3993 8
3994 }
3995 }
3996
3997 let oracle = OracleTopicEmbedder;
3998
3999 let setup = |conn: &rusqlite::Connection| {
4001 for (mid, text) in [
4004 ("m_dup1", "TOPIC_A prefer ripgrep for searching"),
4005 ("m_dup2", "TOPIC_A rg is the fastest searcher"),
4006 ("m_dup3", "TOPIC_A use rg tool to find patterns"),
4007 (
4009 "m_relevant",
4010 "TOPIC_A critical lesson about search performance",
4011 ),
4012 ("m_noise1", "chocolate cookie baking TOPIC_B recipe"),
4014 ("m_noise2", "gardening tulip planting TOPIC_B spring"),
4015 ] {
4016 insert_memory_with_embedding(conn, mid, text, &oracle);
4017 }
4018 };
4019
4020 let weights = kimetsu_core::config::BrokerWeights::default();
4021
4022 let conn_lean = rusqlite::Connection::open_in_memory().expect("in-memory lean");
4029 crate::schema::initialize(&conn_lean).expect("init schema lean");
4030 setup(&conn_lean);
4031
4032 let bundle_lean = retrieve_context_with_embedder(
4033 &conn_lean,
4034 "/fake-repo",
4035 &weights,
4036 ContextRequest {
4037 stage: "localization".to_string(),
4038 query: "TOPIC_A search performance".to_string(),
4039 budget_tokens: 20_000,
4040 min_semantic_score: 0.0, ..Default::default()
4042 },
4043 &[],
4044 &embeddings::NoopEmbedder,
4045 )
4046 .expect("retrieve lean");
4047
4048 let lean_count = bundle_lean
4049 .capsules
4050 .iter()
4051 .filter(|c| c.expansion_handle.starts_with("memory:"))
4052 .count();
4053
4054 let conn_emb = rusqlite::Connection::open_in_memory().expect("in-memory emb");
4056 crate::schema::initialize(&conn_emb).expect("init schema emb");
4057 setup(&conn_emb);
4058
4059 let bundle_emb = retrieve_context_with_embedder(
4060 &conn_emb,
4061 "/fake-repo",
4062 &weights,
4063 ContextRequest {
4064 stage: "localization".to_string(),
4065 query: "TOPIC_A search performance".to_string(),
4066 budget_tokens: 20_000,
4067 min_semantic_score: 0.5, ..Default::default()
4069 },
4070 &[],
4071 &oracle,
4072 )
4073 .expect("retrieve with embeddings");
4074
4075 let emb_count = bundle_emb
4076 .capsules
4077 .iter()
4078 .filter(|c| c.expansion_handle.starts_with("memory:"))
4079 .count();
4080
4081 assert!(
4083 emb_count < lean_count,
4084 "D1e must reduce capsule count: embedding path {emb_count} must be \
4085 < lean path {lean_count}. Embedding capsules: {:?}",
4086 bundle_emb
4087 .capsules
4088 .iter()
4089 .map(|c| &c.expansion_handle)
4090 .collect::<Vec<_>>()
4091 );
4092
4093 assert!(
4095 bundle_emb
4096 .capsules
4097 .iter()
4098 .any(|c| c.expansion_handle == "memory:m_relevant"),
4099 "m_relevant must survive D1e selection (signal preserved); \
4100 embedding capsules: {:?}",
4101 bundle_emb
4102 .capsules
4103 .iter()
4104 .map(|c| &c.expansion_handle)
4105 .collect::<Vec<_>>()
4106 );
4107
4108 let lean_tokens: u32 = bundle_lean.capsules.iter().map(|c| c.token_estimate).sum();
4110 let emb_tokens: u32 = bundle_emb.capsules.iter().map(|c| c.token_estimate).sum();
4111 assert!(
4112 emb_tokens < lean_tokens,
4113 "D1e must reduce token usage: emb={emb_tokens} must be < lean={lean_tokens}"
4114 );
4115 }
4116
4117 #[test]
4123 fn lean_noop_embedder_uses_fts_then_recency_unchanged() {
4124 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4127 crate::schema::initialize(&conn).expect("init schema");
4128
4129 for (mid, text) in [
4131 ("m_x", "use git rebase to clean history"),
4132 ("m_y", "grep finds text quickly"),
4133 ] {
4134 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4135 conn.execute(
4136 "INSERT INTO memories (
4137 memory_id, scope, kind, text, normalized_text, confidence,
4138 source_event_id, provenance_snapshot_json, created_at,
4139 use_count, usefulness_score
4140 )
4141 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
4142 '2026-01-01T00:00:00Z', 0, 0.0)",
4143 rusqlite::params![mid, text, normalized],
4144 )
4145 .expect("insert");
4146 conn.execute(
4147 "INSERT INTO memories_fts (memory_id, text, kind, scope) VALUES (?1, ?2, 'fact', 'global_user')",
4148 rusqlite::params![mid, text],
4149 )
4150 .expect("insert fts");
4151 }
4152
4153 let weights = kimetsu_core::config::BrokerWeights::default();
4154 let bundle = retrieve_context_with_embedder(
4156 &conn,
4157 "/fake-repo",
4158 &weights,
4159 ContextRequest {
4160 stage: "localization".to_string(),
4161 query: "grep text".to_string(),
4162 budget_tokens: 4000,
4163 ..Default::default()
4164 },
4165 &[],
4166 &embeddings::NoopEmbedder,
4167 )
4168 .expect("retrieve with NoopEmbedder must not panic");
4169
4170 let handles: Vec<&str> = bundle
4172 .capsules
4173 .iter()
4174 .filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
4175 .collect();
4176 assert!(
4177 handles.contains(&"m_y"),
4178 "m_y must surface via FTS on lean path; got {handles:?}"
4179 );
4180 }
4182
4183 #[test]
4189 fn classify_task_maps_each_kind_deterministically() {
4190 assert_eq!(
4192 classify_task("fix the panic in the parser"),
4193 TaskKind::Debug,
4194 "contains 'fix' and 'panic'"
4195 );
4196 assert_eq!(
4197 classify_task("there is a crash in auth when calling login"),
4198 TaskKind::Debug,
4199 "contains 'crash'"
4200 );
4201 assert_eq!(
4202 classify_task("debug the failing test"),
4203 TaskKind::Debug,
4204 "contains 'debug' and 'fail'"
4205 );
4206
4207 assert_eq!(
4209 classify_task("investigate why retrieval is slow"),
4210 TaskKind::Investigation,
4211 "contains 'investigate' and 'why'"
4212 );
4213 assert_eq!(
4214 classify_task("analyze the root cause of the latency"),
4215 TaskKind::Investigation,
4216 "contains 'analyze' and 'root cause'"
4217 );
4218
4219 assert_eq!(
4221 classify_task("refactor the auth module"),
4222 TaskKind::Refactor,
4223 "contains 'refactor'"
4224 );
4225 assert_eq!(
4226 classify_task("rename the config struct"),
4227 TaskKind::Refactor,
4228 "contains 'rename'"
4229 );
4230 assert_eq!(
4231 classify_task("simplify the retry handling logic"),
4232 TaskKind::Refactor,
4233 "contains 'simplify'"
4234 );
4235
4236 assert_eq!(
4238 classify_task("document the API endpoints"),
4239 TaskKind::Docs,
4240 "contains 'document'"
4241 );
4242 assert_eq!(
4243 classify_task("update the readme with new instructions"),
4244 TaskKind::Docs,
4245 "contains 'readme'"
4246 );
4247 assert_eq!(
4248 classify_task("add a docstring to the main function"),
4249 TaskKind::Docs,
4250 "contains 'docstring'"
4251 );
4252
4253 assert_eq!(
4255 classify_task("add a dark mode toggle"),
4256 TaskKind::Feature,
4257 "no debug/refactor/docs/investigate keyword"
4258 );
4259 assert_eq!(
4260 classify_task("implement the new caching layer"),
4261 TaskKind::Feature,
4262 "no debug/refactor/docs/investigate keyword"
4263 );
4264 assert_eq!(
4265 classify_task("build the export pipeline"),
4266 TaskKind::Feature,
4267 "no debug/refactor/docs/investigate keyword"
4268 );
4269 }
4270
4271 #[test]
4273 fn classify_task_respects_precedence_order() {
4274 assert_eq!(
4276 classify_task("fix and refactor the login module"),
4277 TaskKind::Debug,
4278 "Debug > Refactor"
4279 );
4280 assert_eq!(
4282 classify_task("investigate and refactor the cache layer"),
4283 TaskKind::Investigation,
4284 "Investigation > Refactor"
4285 );
4286 assert_eq!(
4288 classify_task("investigate the docs and document the API"),
4289 TaskKind::Investigation,
4290 "Investigation > Docs"
4291 );
4292 assert_eq!(
4294 classify_task("refactor and add docs"),
4295 TaskKind::Refactor,
4296 "Refactor > Docs"
4297 );
4298 assert_eq!(
4300 classify_task("fix the bug and investigate the regression"),
4301 TaskKind::Debug,
4302 "Debug > Investigation"
4303 );
4304 }
4305
4306 #[test]
4309 fn weights_for_task_kind_renormalizes_to_unit_sum() {
4310 let base = StageWeights {
4311 relevance: 0.50,
4312 confidence: 0.20,
4313 freshness: 0.20,
4314 scope: 0.10,
4315 };
4316 let original_sum = base.relevance + base.confidence + base.freshness + base.scope;
4317
4318 for kind in [
4319 TaskKind::Debug,
4320 TaskKind::Refactor,
4321 TaskKind::Investigation,
4322 TaskKind::Docs,
4323 ] {
4324 let w = weights_for_task_kind(base.clone(), kind);
4325 let new_sum = w.relevance + w.confidence + w.freshness + w.scope;
4326 assert!(
4328 (new_sum - original_sum).abs() < 1e-4,
4329 "weights_for_task_kind({kind:?}) sum {new_sum} differs from {original_sum}"
4330 );
4331 }
4332 }
4333
4334 #[test]
4336 fn weights_for_task_kind_feature_is_unchanged() {
4337 let base = StageWeights {
4338 relevance: 0.40,
4339 confidence: 0.30,
4340 freshness: 0.20,
4341 scope: 0.10,
4342 };
4343 let w = weights_for_task_kind(base.clone(), TaskKind::Feature);
4344 assert!((w.relevance - base.relevance).abs() < f32::EPSILON);
4345 assert!((w.confidence - base.confidence).abs() < f32::EPSILON);
4346 assert!((w.freshness - base.freshness).abs() < f32::EPSILON);
4347 assert!((w.scope - base.scope).abs() < f32::EPSILON);
4348 }
4349
4350 #[test]
4353 fn weights_for_task_kind_debug_up_freshness_fraction() {
4354 let base = StageWeights {
4355 relevance: 0.50,
4356 confidence: 0.20,
4357 freshness: 0.20,
4358 scope: 0.10,
4359 };
4360 let debug_w = weights_for_task_kind(base.clone(), TaskKind::Debug);
4361 assert!(
4363 debug_w.freshness > base.freshness,
4364 "Debug must increase freshness fraction: {debug_w:?}"
4365 );
4366 }
4367
4368 #[test]
4371 fn weights_for_task_kind_refactor_up_scope_fraction() {
4372 let base = StageWeights {
4373 relevance: 0.50,
4374 confidence: 0.20,
4375 freshness: 0.20,
4376 scope: 0.10,
4377 };
4378 let refactor_w = weights_for_task_kind(base.clone(), TaskKind::Refactor);
4379 assert!(
4380 refactor_w.scope > base.scope,
4381 "Refactor must increase scope fraction: {refactor_w:?}"
4382 );
4383 }
4384
4385 #[test]
4388 fn task_kind_feature_is_retrieval_neutral() {
4389 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4390 crate::schema::initialize(&conn).expect("init schema");
4391
4392 for (mid, db_kind, text) in [
4396 ("m1", "failure_pattern", "linker not found error in build"),
4397 ("m2", "convention", "use snake_case for all identifiers"),
4398 ("m3", "fact", "the cache is invalidated on every deploy"),
4399 ] {
4400 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4401 conn.execute(
4402 "INSERT INTO memories (
4403 memory_id, scope, kind, text, normalized_text, confidence,
4404 source_event_id, provenance_snapshot_json, created_at,
4405 use_count, usefulness_score
4406 )
4407 VALUES (?1, 'project', ?2, ?3, ?4, 1.0, NULL, '{}',
4408 '2026-01-01T00:00:00Z', 0, 0.0)",
4409 rusqlite::params![mid, db_kind, text, normalized],
4410 )
4411 .expect("insert memory");
4412 conn.execute(
4413 "INSERT INTO memories_fts (memory_id, text, kind, scope)
4414 VALUES (?1, ?2, ?3, 'project')",
4415 rusqlite::params![mid, text, db_kind],
4416 )
4417 .expect("insert fts");
4418 }
4419
4420 let weights = kimetsu_core::config::BrokerWeights::default();
4421 let query = "cache convention failure".to_string();
4422
4423 let baseline = retrieve_context_with_embedder(
4425 &conn,
4426 "/fake-repo",
4427 &weights,
4428 ContextRequest {
4429 stage: "localization".to_string(),
4430 query: query.clone(),
4431 budget_tokens: 4000,
4432 ..Default::default()
4433 },
4434 &[],
4435 &embeddings::NoopEmbedder,
4436 )
4437 .expect("baseline retrieve");
4438
4439 let feature = retrieve_context_with_embedder(
4441 &conn,
4442 "/fake-repo",
4443 &weights,
4444 ContextRequest {
4445 stage: "localization".to_string(),
4446 query: query.clone(),
4447 budget_tokens: 4000,
4448 task_kind: TaskKind::Feature,
4449 ..Default::default()
4450 },
4451 &[],
4452 &embeddings::NoopEmbedder,
4453 )
4454 .expect("feature retrieve");
4455
4456 let baseline_ids: Vec<&str> = baseline
4457 .capsules
4458 .iter()
4459 .map(|c| c.expansion_handle.as_str())
4460 .collect();
4461 let feature_ids: Vec<&str> = feature
4462 .capsules
4463 .iter()
4464 .map(|c| c.expansion_handle.as_str())
4465 .collect();
4466 assert_eq!(
4467 baseline_ids, feature_ids,
4468 "task_kind=Feature must produce identical retrieval to default; \
4469 baseline={baseline_ids:?} feature={feature_ids:?}"
4470 );
4471
4472 let baseline_scores: Vec<f32> = baseline.capsules.iter().map(|c| c.score).collect();
4473 let feature_scores: Vec<f32> = feature.capsules.iter().map(|c| c.score).collect();
4474 for (b, f) in baseline_scores.iter().zip(feature_scores.iter()) {
4475 assert!(
4476 (b - f).abs() < 1e-5,
4477 "scores must be identical: baseline={b} feature={f}"
4478 );
4479 }
4480 }
4481
4482 #[test]
4493 fn debug_surfaces_more_failure_pattern_than_docs() {
4494 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4495 crate::schema::initialize(&conn).expect("init schema");
4496
4497 for (i, text) in [
4501 "auth token expired causes login failure",
4502 "auth service crash on null pointer",
4503 "auth regression after upgrade breaks sessions",
4504 "auth error when certificate is invalid",
4505 ]
4506 .iter()
4507 .enumerate()
4508 {
4509 let mid = format!("mfp{i}");
4510 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4511 conn.execute(
4512 "INSERT INTO memories (
4513 memory_id, scope, kind, text, normalized_text, confidence,
4514 source_event_id, provenance_snapshot_json, created_at,
4515 use_count, usefulness_score
4516 )
4517 VALUES (?1, 'project', 'failure_pattern', ?2, ?3, 1.0, NULL, '{}',
4518 '2026-01-01T00:00:00Z', 0, 0.0)",
4519 rusqlite::params![mid, text, normalized],
4520 )
4521 .expect("insert failure_pattern");
4522 conn.execute(
4523 "INSERT INTO memories_fts (memory_id, text, kind, scope)
4524 VALUES (?1, ?2, 'failure_pattern', 'project')",
4525 rusqlite::params![mid, text],
4526 )
4527 .expect("insert fts");
4528 }
4529
4530 for (i, (db_kind, text)) in [
4533 ("convention", "auth module uses bearer tokens by convention"),
4534 ("convention", "auth scopes are documented in the API guide"),
4535 ("fact", "auth service runs on port 8443 in production"),
4536 ("fact", "auth uses JWT with RS256 signing for all tokens"),
4537 ]
4538 .iter()
4539 .enumerate()
4540 {
4541 let mid = format!("mconv{i}");
4542 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4543 conn.execute(
4544 "INSERT INTO memories (
4545 memory_id, scope, kind, text, normalized_text, confidence,
4546 source_event_id, provenance_snapshot_json, created_at,
4547 use_count, usefulness_score
4548 )
4549 VALUES (?1, 'project', ?2, ?3, ?4, 1.0, NULL, '{}',
4550 '2026-01-01T00:00:00Z', 0, 0.0)",
4551 rusqlite::params![mid, db_kind, text, normalized],
4552 )
4553 .expect("insert convention/fact");
4554 conn.execute(
4555 "INSERT INTO memories_fts (memory_id, text, kind, scope)
4556 VALUES (?1, ?2, ?3, 'project')",
4557 rusqlite::params![mid, text, db_kind],
4558 )
4559 .expect("insert fts");
4560 }
4561
4562 let weights = kimetsu_core::config::BrokerWeights::default();
4563 let query = "auth token failure".to_string();
4564
4565 let debug_bundle = retrieve_context_with_embedder(
4567 &conn,
4568 "/fake-repo",
4569 &weights,
4570 ContextRequest {
4571 stage: "localization".to_string(),
4572 query: query.clone(),
4573 budget_tokens: 4000,
4574 max_capsules: 4,
4575 task_kind: TaskKind::Debug,
4576 ..Default::default()
4577 },
4578 &[],
4579 &embeddings::NoopEmbedder,
4580 )
4581 .expect("debug retrieve");
4582
4583 let docs_bundle = retrieve_context_with_embedder(
4585 &conn,
4586 "/fake-repo",
4587 &weights,
4588 ContextRequest {
4589 stage: "localization".to_string(),
4590 query: query.clone(),
4591 budget_tokens: 4000,
4592 max_capsules: 4,
4593 task_kind: TaskKind::Docs,
4594 ..Default::default()
4595 },
4596 &[],
4597 &embeddings::NoopEmbedder,
4598 )
4599 .expect("docs retrieve");
4600
4601 let count_failure_pattern = |bundle: &ContextBundle| -> usize {
4604 bundle
4605 .capsules
4606 .iter()
4607 .filter(|c| capsule_matches_kind(c, "failure_pattern"))
4608 .count()
4609 };
4610
4611 let debug_fp = count_failure_pattern(&debug_bundle);
4612 let docs_fp = count_failure_pattern(&docs_bundle);
4613
4614 assert!(
4615 debug_fp > docs_fp,
4616 "Debug must surface strictly more failure_pattern capsules than Docs: \
4617 debug_fp={debug_fp} docs_fp={docs_fp}\n\
4618 Debug capsules: {:?}\n\
4619 Docs capsules: {:?}",
4620 debug_bundle
4621 .capsules
4622 .iter()
4623 .map(|c| format!("{}:{}", c.kind, &c.summary[..c.summary.len().min(60)]))
4624 .collect::<Vec<_>>(),
4625 docs_bundle
4626 .capsules
4627 .iter()
4628 .map(|c| format!("{}:{}", c.kind, &c.summary[..c.summary.len().min(60)]))
4629 .collect::<Vec<_>>(),
4630 );
4631 }
4632
4633 fn init_db_with_memory(memory_id: &str, text: &str) -> rusqlite::Connection {
4636 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4637 crate::schema::initialize(&conn).expect("init schema");
4638 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4639 conn.execute(
4640 "INSERT INTO memories (
4641 memory_id, scope, kind, text, normalized_text, confidence,
4642 source_event_id, provenance_snapshot_json, created_at,
4643 use_count, usefulness_score
4644 )
4645 VALUES (?1, 'project', 'fact', ?2, ?3, 1.0, NULL, '{}',
4646 '2026-01-01T00:00:00Z', 0, 0.0)",
4647 rusqlite::params![memory_id, text, normalized],
4648 )
4649 .expect("insert memory");
4650 conn
4651 }
4652
4653 #[test]
4655 fn resolve_capsule_memory_returns_full_text() {
4656 let conn = init_db_with_memory("test-mem-id", "Use rg over grep for speed");
4657 let repo_root = std::path::Path::new("/fake-repo");
4658 let result =
4659 resolve_capsule(&conn, repo_root, "memory:test-mem-id").expect("should resolve");
4660 assert_eq!(result, "Use rg over grep for speed");
4661 }
4662
4663 #[test]
4665 fn resolve_capsule_memory_missing_id_returns_err() {
4666 let conn = init_db_with_memory("real-id", "some text");
4667 let repo_root = std::path::Path::new("/fake-repo");
4668 let err = resolve_capsule(&conn, repo_root, "memory:nonexistent-id")
4669 .expect_err("should error for missing memory");
4670 assert!(
4671 err.to_string().contains("no active memory"),
4672 "error message should mention missing: {err}"
4673 );
4674 }
4675
4676 #[test]
4678 fn resolve_capsule_file_returns_bounded_content() {
4679 let dir = make_test_dir("f2_file_resolve");
4680 let content = "hello from the file\n";
4681 std::fs::write(dir.join("notes.txt"), content).expect("write");
4682 let result = resolve_capsule(
4683 &rusqlite::Connection::open_in_memory().expect("open"),
4685 &dir,
4686 "file:notes.txt",
4687 )
4688 .expect("should resolve file");
4689 assert!(result.contains("hello from the file"));
4690 std::fs::remove_dir_all(&dir).ok();
4691 }
4692
4693 #[test]
4695 fn resolve_capsule_file_caps_large_file() {
4696 let dir = make_test_dir("f2_file_cap");
4697 let big = "A".repeat(FILE_EXPAND_CAP_BYTES * 3);
4698 std::fs::write(dir.join("big.txt"), &big).expect("write");
4699 let result = resolve_capsule(
4700 &rusqlite::Connection::open_in_memory().expect("open"),
4701 &dir,
4702 "file:big.txt",
4703 )
4704 .expect("should resolve large file");
4705 assert!(
4706 result.len() <= FILE_EXPAND_CAP_BYTES + 200,
4707 "result should be bounded: got {} bytes",
4708 result.len()
4709 );
4710 assert!(
4711 result.contains("truncated"),
4712 "truncation marker should be present"
4713 );
4714 std::fs::remove_dir_all(&dir).ok();
4715 }
4716
4717 #[test]
4719 fn resolve_capsule_unknown_handle_returns_err() {
4720 let conn = rusqlite::Connection::open_in_memory().expect("open");
4721 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "blob:abc123")
4722 .expect_err("should error");
4723 assert!(
4724 err.to_string().contains("unrecognised handle"),
4725 "got: {err}"
4726 );
4727 }
4728
4729 #[test]
4731 fn resolve_capsule_malformed_handle_returns_err() {
4732 let conn = rusqlite::Connection::open_in_memory().expect("open");
4733 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "justnocolon")
4734 .expect_err("should error");
4735 assert!(
4736 err.to_string().contains("unrecognised handle"),
4737 "got: {err}"
4738 );
4739 }
4740
4741 #[test]
4743 fn resolve_capsule_run_handle_returns_deferred_err() {
4744 let conn = rusqlite::Connection::open_in_memory().expect("open");
4745 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "run:some-run-id")
4746 .expect_err("run: should be deferred err");
4747 assert!(err.to_string().contains("not yet supported"), "got: {err}");
4748 }
4749
4750 #[test]
4752 fn resolve_capsule_file_rejects_absolute_path() {
4753 let conn = rusqlite::Connection::open_in_memory().expect("open");
4754 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "file:/etc/passwd")
4755 .expect_err("should reject absolute path");
4756 assert!(err.to_string().contains("absolute path"), "got: {err}");
4757 }
4758
4759 fn make_capsule(summary: &str, score: f32) -> ContextCapsule {
4762 ContextCapsule {
4763 id: new_id().to_string(),
4764 kind: "memory".to_string(),
4765 summary: summary.to_string(),
4766 token_estimate: 10,
4767 expansion_handle: format!("memory:{}", new_id()),
4768 provenance: vec![],
4769 confidence: 1.0,
4770 freshness: 1.0,
4771 relevance: 1.0,
4772 scope_weight: 1.0,
4773 score,
4774 }
4775 }
4776
4777 #[test]
4780 fn rerank_capsules_reorders_by_query_overlap() {
4781 use crate::embeddings::StubReranker;
4782
4783 let query = "rust async tokio";
4786 let high_overlap = make_capsule("rust async tokio runtime", 0.0);
4787 let low_overlap = make_capsule("python django framework", 0.0);
4788 let capsules = vec![low_overlap.clone(), high_overlap.clone()];
4790
4791 let ranked = rerank_capsules(query, capsules, &StubReranker, 0.0, 0);
4792
4793 assert_eq!(ranked.len(), 2, "both capsules should survive (floor=0)");
4794 assert!(
4796 ranked[0].summary.contains("rust"),
4797 "rust capsule must be first, got: {:?}",
4798 ranked[0].summary
4799 );
4800 assert!(
4802 ranked[0].score > 0.05,
4803 "score must be overwritten by reranker: {}",
4804 ranked[0].score
4805 );
4806 assert!(
4808 ranked[0].score > ranked[1].score,
4809 "high overlap must score higher: {} vs {}",
4810 ranked[0].score,
4811 ranked[1].score
4812 );
4813 }
4814
4815 #[test]
4819 fn rerank_capsules_floor_drops_zero_overlap() {
4820 use crate::embeddings::StubReranker;
4821
4822 let query = "rust async tokio";
4823 let high = make_capsule("rust async tokio runtime", 0.0);
4824 let zero = make_capsule("completely unrelated document xyz", 0.0); let capsules = vec![high, zero];
4827 let ranked = rerank_capsules(query, capsules, &StubReranker, 0.3, 0);
4828
4829 assert_eq!(ranked.len(), 1, "zero-overlap capsule must be dropped");
4831 assert!(
4832 ranked[0].summary.contains("rust"),
4833 "only rust capsule should survive"
4834 );
4835 }
4836
4837 #[test]
4839 fn rerank_capsules_cap_truncates() {
4840 use crate::embeddings::StubReranker;
4841
4842 let query = "alpha beta gamma";
4843 let capsules = vec![
4844 make_capsule("alpha beta gamma delta", 0.0),
4845 make_capsule("alpha beta", 0.0),
4846 make_capsule("alpha", 0.0),
4847 make_capsule("unrelated xyz", 0.0),
4848 ];
4849
4850 let ranked = rerank_capsules(query, capsules, &StubReranker, 0.0, 2);
4851 assert_eq!(ranked.len(), 2, "cap=2 must truncate to 2 results");
4852 assert!(
4854 ranked[0].score >= ranked[1].score,
4855 "results must be sorted descending"
4856 );
4857 }
4858
4859 #[test]
4861 fn rerank_capsules_fail_open_preserves_input_order() {
4862 struct FailingReranker;
4863 impl crate::embeddings::Reranker for FailingReranker {
4864 fn rerank(
4865 &self,
4866 _query: &str,
4867 _docs: &[&str],
4868 ) -> Result<Vec<f32>, crate::embeddings::EmbedderError> {
4869 Err(crate::embeddings::EmbedderError::EmbedFailed(
4870 "simulated failure".into(),
4871 ))
4872 }
4873 fn model_id(&self) -> &str {
4874 "fail-reranker"
4875 }
4876 }
4877
4878 let query = "anything";
4879 let c1 = make_capsule("first capsule", 0.9);
4880 let c2 = make_capsule("second capsule", 0.5);
4881 let c3 = make_capsule("third capsule", 0.1);
4882 let capsules = vec![c1.clone(), c2.clone(), c3.clone()];
4883
4884 let out = rerank_capsules(query, capsules, &FailingReranker, 0.0, 0);
4885
4886 assert_eq!(out.len(), 3, "all capsules must be returned on error");
4888 assert_eq!(out[0].summary, c1.summary, "order must be preserved");
4889 assert_eq!(out[1].summary, c2.summary, "order must be preserved");
4890 assert_eq!(out[2].summary, c3.summary, "order must be preserved");
4891 }
4892
4893 #[test]
4895 fn rerank_capsules_empty_input_returns_empty() {
4896 use crate::embeddings::StubReranker;
4897 let out = rerank_capsules("query", vec![], &StubReranker, 0.0, 0);
4898 assert!(out.is_empty());
4899 }
4900
4901 #[test]
4905 fn compress_for_render_short_text_unchanged() {
4906 let text = "project:fact - Use cargo fmt before committing.";
4907 let out = compress_for_render(text, 3);
4908 assert_eq!(out, text, "short text must not be altered");
4909 }
4910
4911 #[test]
4913 fn compress_for_render_strips_tags_prefix() {
4914 let text = "[tags: rust, cargo] Always run cargo clippy before submitting a PR.";
4915 let out = compress_for_render(text, 3);
4916 assert!(
4917 !out.starts_with('['),
4918 "tags prefix must be stripped, got: {out:?}"
4919 );
4920 assert!(
4921 out.contains("cargo clippy"),
4922 "body must remain, got: {out:?}"
4923 );
4924 }
4925
4926 #[test]
4928 fn compress_for_render_strips_context_suffix() {
4929 let text =
4930 "project:fact - Use cargo fmt. Always clippy clean. (context: Kimetsu brain lesson)";
4931 let out = compress_for_render(text, 5);
4932 assert!(
4933 !out.contains("(context:"),
4934 "context suffix must be stripped, got: {out:?}"
4935 );
4936 assert!(out.contains("cargo fmt"), "body must remain, got: {out:?}");
4937 }
4938
4939 #[test]
4941 fn compress_for_render_caps_sentences() {
4942 let text =
4943 "project:fact - First sentence. Second sentence. Third sentence. Fourth sentence.";
4944 let out = compress_for_render(text, 2);
4945 assert!(out.contains("First"), "first sentence must be present");
4947 assert!(out.contains("Second"), "second sentence must be present");
4948 assert!(
4949 !out.contains("Third"),
4950 "third sentence must be truncated, got: {out:?}"
4951 );
4952 }
4953
4954 #[test]
4956 fn compress_for_render_preserves_scope_prefix() {
4957 let text = "global_user:convention - First rule. Second rule. Third rule. Fourth rule.";
4958 let out = compress_for_render(text, 2);
4959 assert!(
4960 out.starts_with("global_user:convention - "),
4961 "scope prefix must be preserved, got: {out:?}"
4962 );
4963 assert!(out.contains("First"), "first sentence must remain");
4964 assert!(!out.contains("Third"), "third sentence must be truncated");
4965 }
4966
4967 #[test]
4969 fn compress_for_render_empty_input_safe() {
4970 let out = compress_for_render("", 3);
4971 assert_eq!(out, "", "empty input must return empty string");
4972 }
4973
4974 #[test]
4976 fn compress_for_render_zero_max_sentences_returns_original() {
4977 let text = "project:fact - Some lesson that is quite long. It keeps going. And going.";
4978 let out = compress_for_render(text, 0);
4979 assert_eq!(out, text);
4980 }
4981
4982 #[test]
4984 fn compress_for_render_utf8_safe() {
4985 let text = "project:fact - こんにちは世界. Hello world. Third sentence. Fourth sentence.";
4986 let out = compress_for_render(text, 2);
4988 assert!(!out.is_empty(), "UTF-8 text must not produce empty output");
4989 assert!(!out.contains("Third"), "third sentence must be truncated");
4991 }
4992
4993 #[test]
4996 fn compress_for_render_long_memory_reduces_tokens_by_25_percent() {
4997 let long_summary = "project:fact - When a SQLite WAL file exists from a crashed process, \
4999 opening the DB causes the WAL to be replayed. The replayed WAL may contain \
5000 partial writes that corrupt the DB. Always check for WAL files before opening. \
5001 Delete the WAL only after verifying the DB is consistent. Use PRAGMA integrity_check \
5002 to validate after opening. If integrity_check fails, restore from backup. Never \
5003 truncate the WAL without replaying it first. This pattern applies to any \
5004 crash-recovery scenario.";
5005
5006 let raw_tokens = estimate_tokens(long_summary);
5007 assert!(
5008 raw_tokens > 60,
5009 "test precondition: raw memory must be >60 tokens, got {raw_tokens}"
5010 );
5011
5012 let compressed = compress_for_render(long_summary, 3);
5013 let compressed_tokens = estimate_tokens(&compressed);
5014
5015 let reduction = 1.0 - (compressed_tokens as f64 / raw_tokens as f64);
5016 assert!(
5017 reduction >= 0.25,
5018 "compression must reduce tokens by >=25% on long memories; \
5019 raw={raw_tokens} compressed={compressed_tokens} reduction={reduction:.2}"
5020 );
5021 }
5022}