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 crate::reinforce::apply_query_routing(
484 conn,
485 &request.query,
486 query_embedding.as_ref(),
487 &mut candidates,
488 );
489
490 candidates.extend(repo_file_candidates(conn, repo_root, &request.query, 30)?);
491 candidates.extend(manifest_candidates(conn, repo_root, &request.query)?);
492
493 if !request.kinds.is_empty() {
500 candidates.retain(|c| {
501 request
502 .kinds
503 .iter()
504 .any(|k| capsule_matches_kind(&c.capsule, k))
505 });
506 }
507
508 if request.min_lexical_coverage > 0.0 {
525 let content = content_tokens(&request.query);
526 if !content.is_empty() {
527 let idf = corpus_token_idf(conn, &content)?;
528 let total_idf: f32 = content
529 .iter()
530 .map(|t| idf.get(t).copied().unwrap_or(0.0))
531 .sum();
532 if total_idf > f32::EPSILON {
535 candidates.retain(|c| {
536 if c.capsule.kind != "memory" {
537 return true; }
539 if c.cosine.is_some_and(|cos| cos >= SEMANTIC_KEEP_COSINE) {
542 return true;
543 }
544 weighted_coverage(&content, &idf, &c.capsule.summary)
545 >= request.min_lexical_coverage
546 });
547 }
548 }
549 }
550
551 let stage_weights = weights_for_stage(weights, &request.stage);
554 let effective_weights = weights_for_task_kind(stage_weights, request.task_kind);
555 normalize_and_score(&mut candidates, effective_weights);
556
557 let kind_role_hints = task_kind_prefer_roles(request.task_kind);
560 let mut effective_prefer_roles: Vec<String> = request.prefer_roles.clone();
561 for &hint in kind_role_hints {
562 let hint_s = hint.to_string();
563 if !effective_prefer_roles.contains(&hint_s) {
564 effective_prefer_roles.push(hint_s);
565 }
566 }
567
568 if !request.tags.is_empty() || !effective_prefer_roles.is_empty() {
580 let tags_lc: Vec<String> = request
581 .tags
582 .iter()
583 .map(|t| t.to_ascii_lowercase())
584 .collect();
585 for c in &mut candidates {
586 let summary_lc = c.capsule.summary.to_ascii_lowercase();
587 if !tags_lc.is_empty() && tags_lc.iter().any(|t| summary_lc.contains(t.as_str())) {
588 c.capsule.score *= 1.4;
589 }
590 if !effective_prefer_roles.is_empty()
591 && effective_prefer_roles.iter().any(|r| {
592 if c.capsule.kind == "memory" {
600 capsule_matches_kind(&c.capsule, r.as_str())
601 } else {
602 c.capsule.kind.contains(r.as_str())
603 }
604 })
605 {
606 c.capsule.score *= 1.3;
607 }
608 }
609 }
610
611 if query_embedding.is_some() && request.min_semantic_score > 0.0 {
626 candidates.retain(|c| {
627 match c.cosine {
630 Some(cos) => cos >= request.min_semantic_score,
631 None => true,
632 }
633 });
634 }
635
636 candidates.sort_by(|a, b| {
650 b.capsule
651 .score
652 .partial_cmp(&a.capsule.score)
653 .unwrap_or(Ordering::Equal)
654 .then_with(|| {
655 b.capsule
656 .freshness
657 .partial_cmp(&a.capsule.freshness)
658 .unwrap_or(Ordering::Equal)
659 })
660 .then_with(|| a.capsule.expansion_handle.cmp(&b.capsule.expansion_handle))
665 });
666
667 let embedding_mmr_ran = query_embedding.is_some() && !candidates.is_empty();
670 let candidates = if embedding_mmr_ran {
671 apply_candidate_mmr_diversity(candidates, 0.7)
672 } else {
673 candidates
674 };
675
676 let mut capsules = candidates
677 .into_iter()
678 .map(|candidate| candidate.capsule)
679 .collect::<Vec<_>>();
680
681 if !embedding_mmr_ran {
684 capsules.sort_by(|left, right| {
685 right
686 .score
687 .partial_cmp(&left.score)
688 .unwrap_or(Ordering::Equal)
689 .then_with(|| {
690 right
691 .freshness
692 .partial_cmp(&left.freshness)
693 .unwrap_or(Ordering::Equal)
694 })
695 .then_with(|| left.expansion_handle.cmp(&right.expansion_handle))
697 });
698 }
699
700 let top_score = capsules.first().map(|c| c.score).unwrap_or(0.0);
703 if request.min_score > 0.0 && top_score < request.min_score {
704 return Ok(ContextBundle {
705 stage: request.stage,
706 budget_tokens: request.budget_tokens,
707 used_tokens: 0,
708 capsules: Vec::new(),
709 excluded: capsules,
710 skipped: true,
711 top_score,
712 });
713 }
714
715 let capsules = apply_mmr_diversity(capsules, 0.7);
723
724 let capsule_budget = request.budget_tokens / 2;
725 let mut used_tokens = 0u32;
726 let mut included = Vec::new();
727 let mut excluded = Vec::new();
728
729 for capsule in capsules {
730 if request.max_capsules > 0 && included.len() >= request.max_capsules {
732 excluded.push(capsule);
733 continue;
734 }
735 if used_tokens.saturating_add(capsule.token_estimate) <= capsule_budget {
736 used_tokens += capsule.token_estimate;
737 included.push(capsule);
738 } else {
739 excluded.push(capsule);
740 }
741 }
742
743 Ok(ContextBundle {
744 stage: request.stage,
745 budget_tokens: request.budget_tokens,
746 used_tokens,
747 capsules: included,
748 excluded,
749 skipped: false,
750 top_score,
751 })
752}
753
754pub fn search_memories_including_expired(
766 conn: &Connection,
767 limit: u32,
768) -> KimetsuResult<Vec<ContextCapsule>> {
769 let mut stmt = conn.prepare_cached(
770 "
771 SELECT memory_id, scope, kind, text, confidence, created_at,
772 use_count, usefulness_score, valid_from, valid_to
773 FROM memories
774 WHERE invalidated_at IS NULL
775 AND superseded_by IS NULL
776 ORDER BY created_at DESC
777 LIMIT ?1
778 ",
779 )?;
780 let rows = stmt.query_map(params![limit], |row| {
781 Ok((
782 row.get::<_, String>(0)?,
783 row.get::<_, String>(1)?,
784 row.get::<_, String>(2)?,
785 row.get::<_, String>(3)?,
786 row.get::<_, f32>(4)?,
787 row.get::<_, String>(5)?,
788 row.get::<_, i64>(6)?,
789 row.get::<_, f64>(7)?,
790 row.get::<_, Option<String>>(8)?,
791 row.get::<_, Option<String>>(9)?,
792 ))
793 })?;
794 let now_utc = OffsetDateTime::now_utc();
795 let now_rfc3339 = now_utc
796 .format(&time::format_description::well_known::Rfc3339)
797 .unwrap_or_default();
798 let mut capsules = Vec::new();
799 for row in rows {
800 let (
801 memory_id,
802 scope,
803 kind,
804 text,
805 confidence,
806 created_at,
807 _use_count,
808 _usefulness,
809 _valid_from,
810 valid_to,
811 ) = row?;
812 let freshness = freshness(&created_at);
813 let scope_weight = scope_weight(&scope);
814 let suffix = if let Some(ref vt) = valid_to {
816 if vt.as_str() < now_rfc3339.as_str() {
817 format!(" [expired valid_to={vt}]")
818 } else {
819 format!(" [valid_to={vt}]")
820 }
821 } else {
822 String::new()
823 };
824 capsules.push(ContextCapsule {
825 id: new_id().to_string(),
826 kind: "memory".to_string(),
827 summary: format!("{scope}:{kind} - {text}{suffix}"),
828 token_estimate: estimate_tokens(&text) + 8,
829 expansion_handle: format!("memory:{memory_id}"),
830 provenance: vec![ProvenanceRef {
831 source: "Memory".to_string(),
832 id: memory_id,
833 excerpt: Some(excerpt(&text)),
834 }],
835 confidence,
836 freshness,
837 relevance: 0.0,
838 scope_weight,
839 score: 0.0,
840 });
841 }
842 Ok(capsules)
843}
844
845pub fn search_repo_files(
846 conn: &Connection,
847 repo_root: &str,
848 query: &str,
849 limit: u32,
850) -> KimetsuResult<Vec<ContextCapsule>> {
851 let candidates = repo_file_candidates(conn, repo_root, query, limit)?;
852 let mut capsules = candidates
853 .into_iter()
854 .map(|mut candidate| {
855 candidate.capsule.relevance = candidate.raw_relevance;
856 candidate.capsule.score = candidate.raw_relevance;
857 candidate.capsule
858 })
859 .collect::<Vec<_>>();
860 capsules.sort_by(|left, right| {
861 right
862 .score
863 .partial_cmp(&left.score)
864 .unwrap_or(Ordering::Equal)
865 .then_with(|| left.expansion_handle.cmp(&right.expansion_handle))
866 });
867 Ok(capsules)
868}
869
870#[cfg(feature = "embeddings")]
882fn memory_ann_candidates(
883 conn: &Connection,
884 qe: &QueryEmbedding,
885 k: u32,
886 query_tokens: &[String],
887 half_life_days: f32,
888) -> KimetsuResult<Vec<Candidate>> {
889 let handle = crate::ann::handle_for_query(conn, qe.vector.len(), &qe.model_id)?;
891 let hits = handle
892 .read()
893 .unwrap_or_else(|p| p.into_inner())
894 .search(&qe.vector, k as usize)?;
895 let knn_rowids: Vec<i64> = hits.into_iter().map(|(rowid, _dist)| rowid).collect();
899 if knn_rowids.is_empty() {
900 return Ok(Vec::new());
901 }
902
903 let placeholders: String = knn_rowids
905 .iter()
906 .enumerate()
907 .map(|(i, _)| format!("?{}", i + 1))
908 .collect::<Vec<_>>()
909 .join(", ");
910 let sql = format!(
911 "SELECT memory_id, scope, kind, text, confidence, created_at,
912 use_count, usefulness_score, embedding, embedding_model,
913 last_useful_at
914 FROM memories
915 WHERE invalidated_at IS NULL
916 AND superseded_by IS NULL
917 AND (valid_to IS NULL OR valid_to > datetime('now'))
918 AND embedding_model = ?{model_param}
919 AND rowid IN ({placeholders})",
920 model_param = knn_rowids.len() + 1
921 );
922 let mut stmt = conn.prepare(&sql)?;
923 let mut params_vec: Vec<&dyn rusqlite::ToSql> = knn_rowids
924 .iter()
925 .map(|n| n as &dyn rusqlite::ToSql)
926 .collect();
927 params_vec.push(&qe.model_id);
928 let rows_iter = stmt.query_map(params_vec.as_slice(), |row| {
929 Ok((
930 row.get::<_, String>(0)?,
931 row.get::<_, String>(1)?,
932 row.get::<_, String>(2)?,
933 row.get::<_, String>(3)?,
934 row.get::<_, f32>(4)?,
935 row.get::<_, String>(5)?,
936 row.get::<_, i64>(6)?,
937 row.get::<_, f64>(7)?,
938 row.get::<_, Option<Vec<u8>>>(8)?,
939 row.get::<_, Option<String>>(9)?,
940 row.get::<_, Option<String>>(10)?,
941 ))
942 })?;
943
944 let mut candidates = Vec::new();
945 for row in rows_iter {
946 let (
947 memory_id,
948 scope,
949 kind,
950 text,
951 confidence,
952 created_at,
953 use_count,
954 usefulness_score,
955 embedding,
956 embedding_model,
957 last_useful_at,
958 ) = row?;
959 let (cosine, row_vec) =
960 compute_cosine_and_vec(Some(qe), embedding.as_deref(), embedding_model.as_deref());
961 if let Some(candidate) = memory_row_to_candidate(
962 query_tokens,
963 memory_id,
964 scope,
965 kind,
966 text,
967 confidence,
968 created_at,
969 use_count,
970 usefulness_score,
971 last_useful_at,
972 half_life_days,
973 None, cosine,
975 row_vec,
976 ) {
977 candidates.push(candidate);
978 }
979 }
980 Ok(candidates)
981}
982
983pub(crate) fn memory_candidates_flat(
989 conn: &Connection,
990 query: &str,
991 query_embedding: Option<&QueryEmbedding>,
992 half_life_days: f32,
993) -> KimetsuResult<Vec<Candidate>> {
994 memory_candidates(conn, query, query_embedding, half_life_days)
995}
996
997fn memory_candidates(
998 conn: &Connection,
999 query: &str,
1000 query_embedding: Option<&QueryEmbedding>,
1001 half_life_days: f32,
1002) -> KimetsuResult<Vec<Candidate>> {
1003 let query_tokens = query_tokens(query);
1004
1005 #[cfg(feature = "embeddings")]
1010 if let Some(qe) = query_embedding {
1011 let fts_candidates = if let Some(fts_query) = fts_query(query) {
1013 memory_fts_candidates(
1014 conn,
1015 &query_tokens,
1016 &fts_query,
1017 80,
1018 Some(qe),
1019 half_life_days,
1020 )?
1021 } else {
1022 Vec::new()
1023 };
1024
1025 let ann_candidates = memory_ann_candidates(conn, qe, 80, &query_tokens, half_life_days)?;
1027
1028 let mut seen: HashMap<String, usize> = HashMap::new();
1033 let mut merged: Vec<Candidate> = Vec::new();
1034
1035 for candidate in fts_candidates.into_iter().chain(ann_candidates) {
1036 let mid = candidate
1038 .capsule
1039 .expansion_handle
1040 .strip_prefix("memory:")
1041 .unwrap_or(&candidate.capsule.expansion_handle)
1042 .to_string();
1043 if let Some(&idx) = seen.get(&mid) {
1044 if candidate.raw_relevance > merged[idx].raw_relevance {
1046 merged[idx] = candidate;
1047 }
1048 } else {
1049 seen.insert(mid, merged.len());
1050 merged.push(candidate);
1051 }
1052 }
1053
1054 return Ok(merged);
1055 }
1056
1057 if let Some(fts_query) = fts_query(query) {
1059 let candidates = memory_fts_candidates(
1060 conn,
1061 &query_tokens,
1062 &fts_query,
1063 80,
1064 query_embedding,
1065 half_life_days,
1066 )?;
1067 if !candidates.is_empty() {
1068 return Ok(candidates);
1069 }
1070 }
1071
1072 latest_memory_candidates(conn, &query_tokens, 200, query_embedding, half_life_days)
1073}
1074
1075fn latest_memory_candidates(
1076 conn: &Connection,
1077 query_tokens: &[String],
1078 limit: u32,
1079 query_embedding: Option<&QueryEmbedding>,
1080 half_life_days: f32,
1081) -> KimetsuResult<Vec<Candidate>> {
1082 let mut stmt = conn.prepare_cached(
1093 "
1094 SELECT memory_id, scope, kind, text, confidence, created_at,
1095 use_count, usefulness_score, embedding, embedding_model,
1096 last_useful_at
1097 FROM memories
1098 WHERE invalidated_at IS NULL
1099 AND superseded_by IS NULL
1100 AND (valid_to IS NULL OR valid_to > datetime('now'))
1101 ORDER BY created_at DESC
1102 LIMIT ?1
1103 ",
1104 )?;
1105
1106 let rows = stmt.query_map(params![limit], |row| {
1107 Ok((
1108 row.get::<_, String>(0)?,
1109 row.get::<_, String>(1)?,
1110 row.get::<_, String>(2)?,
1111 row.get::<_, String>(3)?,
1112 row.get::<_, f32>(4)?,
1113 row.get::<_, String>(5)?,
1114 row.get::<_, i64>(6)?,
1115 row.get::<_, f64>(7)?,
1116 row.get::<_, Option<Vec<u8>>>(8)?,
1117 row.get::<_, Option<String>>(9)?,
1118 row.get::<_, Option<String>>(10)?,
1119 ))
1120 })?;
1121
1122 let mut candidates = Vec::new();
1123 for row in rows {
1124 let (
1125 memory_id,
1126 scope,
1127 kind,
1128 text,
1129 confidence,
1130 created_at,
1131 use_count,
1132 usefulness_score,
1133 embedding,
1134 embedding_model,
1135 last_useful_at,
1136 ) = row?;
1137 let (cosine, row_vec) = compute_cosine_and_vec(
1138 query_embedding,
1139 embedding.as_deref(),
1140 embedding_model.as_deref(),
1141 );
1142 if let Some(candidate) = memory_row_to_candidate(
1143 query_tokens,
1144 memory_id,
1145 scope,
1146 kind,
1147 text,
1148 confidence,
1149 created_at,
1150 use_count,
1151 usefulness_score,
1152 last_useful_at,
1153 half_life_days,
1154 None,
1155 cosine,
1156 row_vec,
1157 ) {
1158 candidates.push(candidate);
1159 }
1160 }
1161 Ok(candidates)
1162}
1163
1164fn memory_fts_candidates(
1165 conn: &Connection,
1166 query_tokens: &[String],
1167 fts_query: &str,
1168 limit: u32,
1169 query_embedding: Option<&QueryEmbedding>,
1170 half_life_days: f32,
1171) -> KimetsuResult<Vec<Candidate>> {
1172 let mut stmt = conn.prepare_cached(
1173 "
1174 SELECT m.memory_id, m.scope, m.kind, m.text, m.confidence, m.created_at,
1175 m.use_count, m.usefulness_score, bm25(memories_fts) AS rank,
1176 m.embedding, m.embedding_model, m.last_useful_at
1177 FROM memories_fts
1178 JOIN memories m
1179 ON m.memory_id = memories_fts.memory_id
1180 WHERE m.invalidated_at IS NULL
1181 AND m.superseded_by IS NULL
1182 AND (m.valid_to IS NULL OR m.valid_to > datetime('now'))
1183 AND memories_fts MATCH ?1
1184 ORDER BY rank
1185 LIMIT ?2
1186 ",
1187 )?;
1188
1189 let rows = stmt.query_map(params![fts_query, limit], |row| {
1190 Ok((
1191 row.get::<_, String>(0)?,
1192 row.get::<_, String>(1)?,
1193 row.get::<_, String>(2)?,
1194 row.get::<_, String>(3)?,
1195 row.get::<_, f32>(4)?,
1196 row.get::<_, String>(5)?,
1197 row.get::<_, i64>(6)?,
1198 row.get::<_, f64>(7)?,
1199 row.get::<_, f64>(8)?,
1200 row.get::<_, Option<Vec<u8>>>(9)?,
1201 row.get::<_, Option<String>>(10)?,
1202 row.get::<_, Option<String>>(11)?,
1203 ))
1204 })?;
1205
1206 let mut candidates = Vec::new();
1207 for row in rows {
1208 let (
1209 memory_id,
1210 scope,
1211 kind,
1212 text,
1213 confidence,
1214 created_at,
1215 use_count,
1216 usefulness_score,
1217 rank,
1218 embedding,
1219 embedding_model,
1220 last_useful_at,
1221 ) = row?;
1222 let fts_relevance = (-rank as f32).max(0.0);
1223 let (cosine, row_vec) = compute_cosine_and_vec(
1224 query_embedding,
1225 embedding.as_deref(),
1226 embedding_model.as_deref(),
1227 );
1228 if let Some(candidate) = memory_row_to_candidate(
1229 query_tokens,
1230 memory_id,
1231 scope,
1232 kind,
1233 text,
1234 confidence,
1235 created_at,
1236 use_count,
1237 usefulness_score,
1238 last_useful_at,
1239 half_life_days,
1240 Some(fts_relevance),
1241 cosine,
1242 row_vec,
1243 ) {
1244 candidates.push(candidate);
1245 }
1246 }
1247 Ok(candidates)
1248}
1249
1250fn compute_cosine_and_vec(
1274 query_embedding: Option<&QueryEmbedding>,
1275 row_bytes: Option<&[u8]>,
1276 row_model: Option<&str>,
1277) -> (Option<f32>, Option<Vec<f32>>) {
1278 let q = match query_embedding {
1279 Some(q) => q,
1280 None => return (None, None),
1281 };
1282 let bytes = match row_bytes {
1283 Some(b) => b,
1284 None => return (None, None),
1285 };
1286 let model = match row_model {
1287 Some(m) => m,
1288 None => return (None, None),
1289 };
1290 if model != q.model_id {
1291 return (None, None);
1292 }
1293 let row_vec = match decode_embedding(bytes, Some(q.vector.len())) {
1294 Ok(v) => v,
1295 Err(_) => return (None, None),
1296 };
1297 let score = cosine_similarity(&q.vector, &row_vec);
1298 (Some(score), Some(row_vec))
1299}
1300
1301#[allow(clippy::too_many_arguments)]
1302fn memory_row_to_candidate(
1303 query_tokens: &[String],
1304 memory_id: String,
1305 scope: String,
1306 kind: String,
1307 text: String,
1308 confidence: f32,
1309 created_at: String,
1310 use_count: i64,
1311 usefulness_score: f64,
1312 last_useful_at: Option<String>,
1313 half_life_days: f32,
1314 raw_relevance_override: Option<f32>,
1315 cosine_score: Option<f32>,
1316 row_embedding: Option<Vec<f32>>,
1321) -> Option<Candidate> {
1322 let lexical = lexical_relevance(query_tokens, &format!("{kind} {text}"));
1323 let lexical_term = raw_relevance_override.unwrap_or(lexical).max(lexical);
1324
1325 let raw_relevance = match cosine_score {
1337 Some(c) => {
1338 let normalized_cos = ((c + 1.0) * 0.5).clamp(0.0, 1.0);
1339 (1.0 - DEFAULT_HYBRID_ALPHA) * lexical_term + DEFAULT_HYBRID_ALPHA * normalized_cos
1340 }
1341 None => lexical_term,
1342 };
1343
1344 if raw_relevance <= 0.0 && !query_tokens.is_empty() {
1350 return None;
1351 }
1352
1353 let freshness = freshness(&created_at);
1354 let scope_weight = scope_weight(&scope);
1355 let raw_multiplier = usefulness_multiplier(usefulness_score as f32, use_count as u32);
1361 let decay = usefulness_decay(last_useful_at.as_deref(), &created_at, half_life_days);
1362 let multiplier = 1.0 + (raw_multiplier - 1.0) * decay;
1363 let biased_relevance = apply_usefulness_boost(raw_relevance, multiplier);
1364 Some(Candidate {
1365 raw_relevance: biased_relevance,
1366 embedding: row_embedding,
1367 cosine: cosine_score,
1368 capsule: ContextCapsule {
1369 id: new_id().to_string(),
1370 kind: "memory".to_string(),
1371 summary: format!("{scope}:{kind} - {text}"),
1372 token_estimate: estimate_tokens(&text) + 8,
1373 expansion_handle: format!("memory:{memory_id}"),
1374 provenance: vec![ProvenanceRef {
1375 source: "Memory".to_string(),
1376 id: memory_id,
1377 excerpt: Some(excerpt(&text)),
1378 }],
1379 confidence,
1380 freshness,
1381 relevance: 0.0,
1382 scope_weight,
1383 score: 0.0,
1384 },
1385 })
1386}
1387
1388pub(crate) fn usefulness_decay(
1412 last_useful_at: Option<&str>,
1413 created_at: &str,
1414 half_life_days: f32,
1415) -> f32 {
1416 if half_life_days <= 0.0 {
1417 return 1.0;
1418 }
1419 let reference = last_useful_at.unwrap_or(created_at);
1420 let Ok(reference_ts) =
1421 OffsetDateTime::parse(reference, &time::format_description::well_known::Rfc3339)
1422 else {
1423 return 1.0;
1424 };
1425 let age = OffsetDateTime::now_utc() - reference_ts;
1426 let age_days = (age.whole_seconds().max(0) as f32) / 86_400.0;
1427 let exponent = -std::f32::consts::LN_2 * age_days / half_life_days;
1428 exponent.exp().clamp(0.0, 1.0)
1429}
1430
1431pub(crate) use crate::scoring::USEFULNESS_BOOST_CAP;
1434
1435pub(crate) fn apply_usefulness_boost(raw_relevance: f32, multiplier: f32) -> f32 {
1438 if multiplier <= 1.0 {
1439 return raw_relevance * multiplier;
1440 }
1441 (raw_relevance * multiplier).min(raw_relevance + USEFULNESS_BOOST_CAP)
1442}
1443
1444pub(crate) fn usefulness_multiplier(usefulness_score: f32, use_count: u32) -> f32 {
1449 use crate::scoring::{FULL_CONFIDENCE_USES, MULTIPLIER_MAX, MULTIPLIER_MIN};
1457 if use_count == 0 {
1458 return 1.0;
1459 }
1460 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);
1463 let confidence = (use_count as f32 / FULL_CONFIDENCE_USES as f32).min(1.0);
1464 1.0 * (1.0 - confidence) + full_multiplier * confidence
1465}
1466
1467fn repo_file_candidates(
1468 conn: &Connection,
1469 repo_root: &str,
1470 query: &str,
1471 limit: u32,
1472) -> KimetsuResult<Vec<Candidate>> {
1473 let Some(fts_query) = fts_query(query) else {
1474 return Ok(Vec::new());
1475 };
1476
1477 let mut stmt = conn.prepare_cached(
1478 "
1479 SELECT path, snippet, language_guess, bm25(repo_files_fts) AS rank
1480 FROM repo_files_fts
1481 WHERE repo_root = ?1 AND repo_files_fts MATCH ?2
1482 ORDER BY rank
1483 LIMIT ?3
1484 ",
1485 )?;
1486
1487 let rows = stmt.query_map(params![repo_root, fts_query, limit], |row| {
1488 Ok((
1489 row.get::<_, String>(0)?,
1490 row.get::<_, String>(1)?,
1491 row.get::<_, String>(2)?,
1492 row.get::<_, f64>(3)?,
1493 ))
1494 })?;
1495
1496 let mut candidates = Vec::new();
1497 for row in rows {
1498 let (path, snippet, language, rank) = row?;
1499 let raw_relevance = (-rank as f32).max(0.0);
1500 let summary = format!("{path} ({language}) - {}", excerpt(&snippet));
1501 let token_estimate = estimate_tokens(&summary) + 8;
1502 candidates.push(Candidate {
1503 raw_relevance,
1504 embedding: None,
1505 cosine: None,
1506 capsule: ContextCapsule {
1507 id: new_id().to_string(),
1508 kind: "repo_file".to_string(),
1509 summary,
1510 token_estimate,
1511 expansion_handle: format!("file:{path}"),
1512 provenance: vec![ProvenanceRef {
1513 source: "RepoFile".to_string(),
1514 id: path.clone(),
1515 excerpt: Some(excerpt(&snippet)),
1516 }],
1517 confidence: 0.9,
1518 freshness: 1.0,
1519 relevance: 0.0,
1520 scope_weight: 0.9,
1521 score: 0.0,
1522 },
1523 });
1524 }
1525 Ok(candidates)
1526}
1527
1528fn manifest_candidates(
1529 conn: &Connection,
1530 repo_root: &str,
1531 query: &str,
1532) -> KimetsuResult<Vec<Candidate>> {
1533 if let Some(fts_query) = fts_query(query) {
1534 let candidates = manifest_fts_candidates(conn, repo_root, &fts_query, 30)?;
1535 if !candidates.is_empty() {
1536 return Ok(candidates);
1537 }
1538 }
1539
1540 let query_tokens = query_tokens(query);
1541 let mut stmt = conn.prepare_cached(
1542 "
1543 SELECT manifest_path, manifest_kind, parsed_summary_json
1544 FROM repo_manifests
1545 WHERE repo_root = ?1
1546 ORDER BY manifest_path
1547 ",
1548 )?;
1549
1550 let rows = stmt.query_map(params![repo_root], |row| {
1551 Ok((
1552 row.get::<_, String>(0)?,
1553 row.get::<_, String>(1)?,
1554 row.get::<_, String>(2)?,
1555 ))
1556 })?;
1557
1558 let mut candidates = Vec::new();
1559 for row in rows {
1560 let (path, kind, summary_json) = row?;
1561 let raw_relevance =
1562 lexical_relevance(&query_tokens, &format!("{path} {kind} {summary_json}"));
1563 if raw_relevance <= 0.0 && !query_tokens.is_empty() {
1564 continue;
1565 }
1566 let summary = format!("{path} manifest ({kind})");
1567 let token_estimate = estimate_tokens(&summary) + 8;
1568 candidates.push(Candidate {
1569 raw_relevance,
1570 embedding: None,
1571 cosine: None,
1572 capsule: ContextCapsule {
1573 id: new_id().to_string(),
1574 kind: "repo_manifest".to_string(),
1575 summary,
1576 token_estimate,
1577 expansion_handle: format!("file:{path}"),
1578 provenance: vec![ProvenanceRef {
1579 source: "Manifest".to_string(),
1580 id: path,
1581 excerpt: Some(excerpt(&summary_json)),
1582 }],
1583 confidence: 0.95,
1584 freshness: 1.0,
1585 relevance: 0.0,
1586 scope_weight: 0.9,
1587 score: 0.0,
1588 },
1589 });
1590 }
1591 Ok(candidates)
1592}
1593
1594fn manifest_fts_candidates(
1595 conn: &Connection,
1596 repo_root: &str,
1597 fts_query: &str,
1598 limit: u32,
1599) -> KimetsuResult<Vec<Candidate>> {
1600 let mut stmt = conn.prepare_cached(
1601 "
1602 SELECT manifest_path, manifest_kind, parsed_summary_json,
1603 bm25(repo_manifests_fts) AS rank
1604 FROM repo_manifests_fts
1605 WHERE repo_root = ?1 AND repo_manifests_fts MATCH ?2
1606 ORDER BY rank
1607 LIMIT ?3
1608 ",
1609 )?;
1610
1611 let rows = stmt.query_map(params![repo_root, fts_query, limit], |row| {
1612 Ok((
1613 row.get::<_, String>(0)?,
1614 row.get::<_, String>(1)?,
1615 row.get::<_, String>(2)?,
1616 row.get::<_, f64>(3)?,
1617 ))
1618 })?;
1619
1620 let mut candidates = Vec::new();
1621 for row in rows {
1622 let (path, kind, summary_json, rank) = row?;
1623 let raw_relevance = (-rank as f32).max(0.0);
1624 let summary = format!("{path} manifest ({kind})");
1625 let token_estimate = estimate_tokens(&summary) + 8;
1626 candidates.push(Candidate {
1627 raw_relevance,
1628 embedding: None,
1629 cosine: None,
1630 capsule: ContextCapsule {
1631 id: new_id().to_string(),
1632 kind: "repo_manifest".to_string(),
1633 summary,
1634 token_estimate,
1635 expansion_handle: format!("file:{path}"),
1636 provenance: vec![ProvenanceRef {
1637 source: "Manifest".to_string(),
1638 id: path,
1639 excerpt: Some(excerpt(&summary_json)),
1640 }],
1641 confidence: 0.95,
1642 freshness: 1.0,
1643 relevance: 0.0,
1644 scope_weight: 0.9,
1645 score: 0.0,
1646 },
1647 });
1648 }
1649 Ok(candidates)
1650}
1651
1652fn normalize_and_score(candidates: &mut [Candidate], weights: StageWeights) {
1653 let mut max_by_kind = HashMap::<String, f32>::new();
1654 for candidate in candidates.iter() {
1655 max_by_kind
1656 .entry(candidate.capsule.kind.clone())
1657 .and_modify(|max| *max = (*max).max(candidate.raw_relevance))
1658 .or_insert(candidate.raw_relevance);
1659 }
1660
1661 for candidate in candidates {
1662 let max = max_by_kind
1663 .get(&candidate.capsule.kind)
1664 .copied()
1665 .unwrap_or(0.0);
1666 let relevance = if max <= f32::EPSILON {
1667 if candidate.raw_relevance > 0.0 {
1668 1.0
1669 } else {
1670 0.0
1671 }
1672 } else {
1673 (candidate.raw_relevance / max).clamp(0.0, 1.0)
1674 };
1675 candidate.capsule.relevance = relevance;
1676 candidate.capsule.score = weights.relevance * relevance
1677 + weights.confidence * candidate.capsule.confidence
1678 + weights.freshness * candidate.capsule.freshness
1679 + weights.scope * candidate.capsule.scope_weight;
1680 }
1681}
1682
1683fn weights_for_stage(weights: &BrokerWeights, stage: &str) -> StageWeights {
1684 match stage {
1685 "localization" => weights.localization.clone(),
1686 "patch_plan" => weights.patch_plan.clone(),
1687 "verification" => weights.verification.clone(),
1688 "review" => weights.review.clone(),
1689 _ => None,
1690 }
1691 .unwrap_or(StageWeights {
1692 relevance: weights.relevance,
1693 confidence: weights.confidence,
1694 freshness: weights.freshness,
1695 scope: weights.scope,
1696 })
1697}
1698
1699pub(crate) fn scope_weight_pub(scope: &str) -> f32 {
1702 scope_weight(scope)
1703}
1704
1705fn scope_weight(scope: &str) -> f32 {
1706 match scope.parse::<MemoryScope>() {
1707 Ok(MemoryScope::Run) => 1.0,
1708 Ok(MemoryScope::Repo) => 0.9,
1709 Ok(MemoryScope::Project) => 0.7,
1710 Ok(MemoryScope::GlobalUser) => 0.5,
1711 Err(_) => 0.3,
1712 }
1713}
1714
1715pub(crate) fn freshness_pub(created_at: &str) -> f32 {
1718 freshness(created_at)
1719}
1720
1721fn freshness(created_at: &str) -> f32 {
1722 let Ok(created_at) =
1723 OffsetDateTime::parse(created_at, &time::format_description::well_known::Rfc3339)
1724 else {
1725 return 0.5;
1726 };
1727 let age = OffsetDateTime::now_utc() - created_at;
1728 let age_days = age.whole_seconds().max(0) as f32 / 86_400.0;
1729 (-age_days / 30.0).exp().clamp(0.0, 1.0)
1730}
1731
1732const SEMANTIC_KEEP_COSINE: f32 = 0.20;
1737
1738const STOPWORDS: &[&str] = &[
1743 "the", "and", "for", "are", "but", "not", "you", "your", "with", "this", "that", "these",
1744 "those", "from", "into", "about", "what", "whats", "which", "who", "whom", "how", "why",
1745 "when", "where", "can", "could", "would", "should", "will", "shall", "does", "did", "was",
1746 "were", "been", "being", "have", "has", "had", "its", "it", "is", "as", "at", "by", "of", "to",
1747 "in", "on", "or", "an", "be", "do", "me", "my", "we", "us", "our", "im", "ive", "let", "lets",
1748 "please", "tell", "give", "show", "want", "need", "get", "got", "use", "using", "there",
1749 "their", "they", "them", "then", "than", "some", "any", "all", "more", "most", "such", "via",
1750 "per",
1751];
1752
1753fn content_tokens(query: &str) -> Vec<String> {
1758 let mut seen = std::collections::HashSet::new();
1759 query
1760 .split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
1761 .map(str::trim)
1762 .filter(|part| part.len() >= 2)
1763 .map(str::to_ascii_lowercase)
1764 .filter(|t| !STOPWORDS.contains(&t.as_str()))
1765 .map(|t| light_stem(&t).to_string())
1768 .filter(|t| seen.insert(t.clone()))
1769 .collect()
1770}
1771
1772fn corpus_token_idf(conn: &Connection, tokens: &[String]) -> KimetsuResult<HashMap<String, f32>> {
1790 let mut idf = HashMap::new();
1791 let n: i64 = conn
1792 .query_row(
1793 "SELECT COUNT(*) FROM memories WHERE invalidated_at IS NULL",
1794 [],
1795 |row| row.get(0),
1796 )
1797 .unwrap_or(0);
1798 if n == 0 {
1799 return Ok(idf);
1800 }
1801 let mut stmt = conn.prepare_cached(
1802 "SELECT COUNT(*) FROM memories \
1803 WHERE invalidated_at IS NULL AND lower(text) LIKE ?1 ESCAPE '\\'",
1804 )?;
1805 for token in tokens {
1806 let pattern = format!("%{}%", escape_like(token));
1807 let df: i64 = stmt
1808 .query_row(params![pattern], |row| row.get(0))
1809 .unwrap_or(0);
1810 let weight = if df == 0 {
1812 0.0
1813 } else {
1814 (((n + 1) as f32) / ((df + 1) as f32)).ln().max(0.0)
1815 };
1816 idf.insert(token.clone(), weight);
1817 }
1818 Ok(idf)
1819}
1820
1821fn escape_like(token: &str) -> String {
1824 token
1825 .replace('\\', "\\\\")
1826 .replace('%', "\\%")
1827 .replace('_', "\\_")
1828}
1829
1830fn weighted_coverage(content: &[String], idf: &HashMap<String, f32>, summary: &str) -> f32 {
1837 let haystack = summary.to_ascii_lowercase();
1838 let mut total = 0.0f32;
1839 let mut hit = 0.0f32;
1840 for token in content {
1841 let weight = idf.get(token).copied().unwrap_or(0.0);
1842 total += weight;
1843 if weight > 0.0 && haystack.contains(token.as_str()) {
1844 hit += weight;
1845 }
1846 }
1847 if total <= f32::EPSILON {
1848 0.0
1849 } else {
1850 (hit / total).clamp(0.0, 1.0)
1851 }
1852}
1853
1854fn light_stem(token: &str) -> &str {
1865 for suffix in ["ing", "ed", "es", "s"] {
1866 if let Some(stem) = token.strip_suffix(suffix)
1867 && stem.len() >= 4
1868 {
1869 return stem;
1870 }
1871 }
1872 token
1873}
1874
1875fn query_tokens(query: &str) -> Vec<String> {
1876 let mut tokens: Vec<String> = query
1877 .split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
1878 .map(str::trim)
1879 .filter(|part| part.len() >= 2)
1880 .map(str::to_ascii_lowercase)
1881 .map(|t| light_stem(&t).to_string())
1882 .collect();
1883 let lower = query.to_ascii_lowercase();
1890 for (triggers, expansions) in CLASS_HINTS.iter() {
1891 if triggers.iter().any(|t| lower.contains(t)) {
1892 tokens.extend(expansions.iter().map(|e| e.to_string()));
1893 }
1894 }
1895 tokens
1896}
1897
1898const CLASS_HINTS: &[(&[&str], &[&str])] = &[
1906 (
1907 &[
1908 "build",
1909 "compile",
1910 "make",
1911 "cargo",
1912 "cmake",
1913 "configure",
1914 "install",
1915 "train",
1916 "benchmark",
1917 "test suite",
1918 "ray trace",
1919 "render",
1920 ],
1921 &[
1922 "shell_background",
1923 "shell_status",
1924 "shell_output",
1925 "shell_stop",
1926 "long_running",
1927 ],
1928 ),
1929 (
1930 &[
1931 "edit", "modify", "change", "fix", "update", "patch", "refactor", "rename",
1932 ],
1933 &["edit_file", "apply_patch", "old_string", "new_string"],
1934 ),
1935 (
1936 &[
1937 "read", "inspect", "review", "analyze", "examine", "view", "show",
1938 ],
1939 &["read_file", "offset", "limit", "multi_read"],
1940 ),
1941 (
1942 &["find", "locate", "search", "look up", "discover", "list"],
1943 &["glob", "search_files", "list_files"],
1944 ),
1945 (
1946 &["plan", "step", "checklist", "todo", "task list", "phase"],
1947 &["plan", "todos"],
1948 ),
1949 (
1950 &[
1951 "verify",
1952 "check",
1953 "ensure",
1954 "validate",
1955 "pass test",
1956 "verifier",
1957 ],
1958 &["finish", "verifier", "verification"],
1959 ),
1960 (
1961 &[
1962 "image",
1963 "png",
1964 "jpeg",
1965 "jpg",
1966 "pdf",
1967 "diagram",
1968 "screenshot",
1969 ],
1970 &["view_image", "base64", "sha256"],
1971 ),
1972 (&["delete", "remove", "rm "], &["delete_file", "recursive"]),
1973 (&["rename", "move file", "mv "], &["move_file"]),
1974];
1975
1976fn capsule_matches_kind(capsule: &ContextCapsule, wanted: &str) -> bool {
1981 if capsule.kind == wanted {
1982 return true;
1983 }
1984 if capsule.kind == "memory"
1985 && let Some((prefix, _)) = capsule.summary.split_once(" - ")
1986 && let Some((_scope, mkind)) = prefix.split_once(':')
1987 {
1988 return mkind == wanted;
1989 }
1990 false
1991}
1992
1993pub(crate) fn fts_query(query: &str) -> Option<String> {
1994 let tokens = query_tokens(query);
1995 if tokens.is_empty() {
1996 return None;
1997 }
1998 Some(
1999 tokens
2000 .into_iter()
2001 .take(12)
2002 .map(|token| format!("{token}*"))
2003 .collect::<Vec<_>>()
2004 .join(" OR "),
2005 )
2006}
2007
2008fn apply_candidate_mmr_diversity(mut sorted: Vec<Candidate>, lambda: f32) -> Vec<Candidate> {
2028 if sorted.len() <= 1 {
2029 return sorted;
2030 }
2031 let summaries: Vec<std::collections::HashSet<String>> = sorted
2033 .iter()
2034 .map(|c| summary_token_set(&c.capsule.summary))
2035 .collect();
2036
2037 let mut picked_indices: Vec<usize> = Vec::with_capacity(sorted.len());
2038 let mut remaining: Vec<usize> = (0..sorted.len()).collect();
2039
2040 picked_indices.push(remaining.remove(0));
2042
2043 while !remaining.is_empty() {
2044 let mut best_idx_in_remaining = 0;
2045 let mut best_score = f32::MIN;
2046
2047 for (i, &cand) in remaining.iter().enumerate() {
2048 let mut max_overlap = 0.0f32;
2049 for &p in &picked_indices {
2050 let same_kind = sorted[cand].capsule.kind == sorted[p].capsule.kind;
2053 let raw_overlap = candidate_pair_overlap(
2054 &sorted[cand],
2055 &sorted[p],
2056 &summaries[cand],
2057 &summaries[p],
2058 );
2059 let overlap = if same_kind {
2060 raw_overlap
2061 } else {
2062 raw_overlap * 0.5
2063 };
2064 if overlap > max_overlap {
2065 max_overlap = overlap;
2066 }
2067 }
2068 let mmr = lambda * sorted[cand].capsule.score - (1.0 - lambda) * max_overlap;
2069 if mmr > best_score {
2070 best_score = mmr;
2071 best_idx_in_remaining = i;
2072 }
2073 }
2074 picked_indices.push(remaining.remove(best_idx_in_remaining));
2075 }
2076
2077 let mut taken: Vec<Option<Candidate>> = sorted.drain(..).map(Some).collect();
2079 let mut out = Vec::with_capacity(taken.len());
2080 for idx in picked_indices {
2081 if let Some(c) = taken[idx].take() {
2082 out.push(c);
2083 }
2084 }
2085 out
2086}
2087
2088fn candidate_pair_overlap(
2095 a: &Candidate,
2096 b: &Candidate,
2097 tokens_a: &std::collections::HashSet<String>,
2098 tokens_b: &std::collections::HashSet<String>,
2099) -> f32 {
2100 if let (Some(va), Some(vb)) = (a.embedding.as_deref(), b.embedding.as_deref()) {
2101 cosine_similarity(va, vb).max(0.0)
2106 } else {
2107 jaccard(tokens_a, tokens_b)
2108 }
2109}
2110
2111fn apply_mmr_diversity(mut sorted: Vec<ContextCapsule>, lambda: f32) -> Vec<ContextCapsule> {
2123 if sorted.len() <= 1 {
2124 return sorted;
2125 }
2126 let summaries: Vec<std::collections::HashSet<String>> = sorted
2128 .iter()
2129 .map(|c| summary_token_set(&c.summary))
2130 .collect();
2131 let mut picked_indices: Vec<usize> = Vec::with_capacity(sorted.len());
2132 let mut remaining: Vec<usize> = (0..sorted.len()).collect();
2133
2134 picked_indices.push(remaining.remove(0));
2136
2137 while !remaining.is_empty() {
2138 let mut best_idx_in_remaining = 0;
2139 let mut best_score = f32::MIN;
2140 for (i, &cand) in remaining.iter().enumerate() {
2141 let mut max_overlap = 0.0f32;
2142 for &p in &picked_indices {
2143 let raw = jaccard(&summaries[cand], &summaries[p]);
2144 let overlap = if sorted[cand].kind == sorted[p].kind {
2145 raw
2146 } else {
2147 raw * 0.5
2150 };
2151 if overlap > max_overlap {
2152 max_overlap = overlap;
2153 }
2154 }
2155 let mmr = lambda * sorted[cand].score - (1.0 - lambda) * max_overlap;
2156 if mmr > best_score {
2157 best_score = mmr;
2158 best_idx_in_remaining = i;
2159 }
2160 }
2161 picked_indices.push(remaining.remove(best_idx_in_remaining));
2162 }
2163 let mut out = Vec::with_capacity(sorted.len());
2165 let mut taken: Vec<Option<ContextCapsule>> = sorted.drain(..).map(Some).collect();
2167 for idx in picked_indices {
2168 if let Some(c) = taken[idx].take() {
2169 out.push(c);
2170 }
2171 }
2172 out
2173}
2174
2175fn summary_token_set(s: &str) -> std::collections::HashSet<String> {
2176 s.split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_')
2177 .filter(|t| t.len() >= 3)
2178 .map(str::to_ascii_lowercase)
2179 .collect()
2180}
2181
2182fn jaccard(a: &std::collections::HashSet<String>, b: &std::collections::HashSet<String>) -> f32 {
2183 if a.is_empty() && b.is_empty() {
2184 return 0.0;
2185 }
2186 let intersection = a.intersection(b).count();
2187 let union = a.union(b).count();
2188 intersection as f32 / union.max(1) as f32
2189}
2190
2191fn lexical_relevance(tokens: &[String], haystack: &str) -> f32 {
2192 if tokens.is_empty() {
2193 return 0.0;
2194 }
2195 let haystack = haystack.to_ascii_lowercase();
2196 let matches = tokens
2197 .iter()
2198 .filter(|token| haystack.contains(token.as_str()))
2199 .count();
2200 matches as f32 / tokens.len() as f32
2201}
2202
2203pub fn estimate_tokens(text: &str) -> u32 {
2204 ((text.split_whitespace().count() as f32) * 1.33).ceil() as u32
2205}
2206
2207pub fn compress_for_render(summary: &str, max_sentences: usize) -> String {
2230 if max_sentences == 0 {
2231 return summary.to_string();
2232 }
2233
2234 let text = if let Some(rest) = summary.strip_prefix('[') {
2236 if let Some(idx) = rest.find(']') {
2238 rest[idx + 1..].trim_start()
2239 } else {
2240 summary
2241 }
2242 } else {
2243 summary
2244 };
2245
2246 let text = if let Some(idx) = text.rfind('(') {
2248 let candidate = text[..idx].trim_end();
2249 let inner = &text[idx + 1..];
2252 if inner.contains(':') && inner.trim_end().ends_with(')') {
2253 candidate
2254 } else {
2255 text
2256 }
2257 } else {
2258 text
2259 };
2260
2261 let (scope_prefix, body) = if let Some(dash_pos) = text.find(" - ") {
2263 let prefix_candidate = &text[..dash_pos];
2264 if !prefix_candidate.contains(' ') && prefix_candidate.contains(':') {
2266 let body_start = dash_pos + 3; (&text[..body_start], &text[body_start..])
2268 } else {
2269 ("", text)
2270 }
2271 } else {
2272 ("", text)
2273 };
2274
2275 let compressed_body = cap_sentences(body, max_sentences);
2277
2278 let result = if scope_prefix.is_empty() {
2280 compressed_body.to_string()
2281 } else {
2282 format!("{scope_prefix}{compressed_body}")
2283 };
2284
2285 if result.trim().is_empty() {
2286 summary.to_string()
2287 } else {
2288 result
2289 }
2290}
2291
2292fn cap_sentences(text: &str, n: usize) -> &str {
2296 let bytes = text.as_bytes();
2297 let len = bytes.len();
2298 let mut count = 0;
2299 let mut i = 0;
2300 while i < len {
2301 if bytes[i] == b'.' {
2303 let next = i + 1;
2304 if next < len && (bytes[next] == b' ' || bytes[next] == b'\n') {
2305 count += 1;
2306 if count >= n {
2307 return text[..=i].trim_end();
2309 }
2310 }
2311 }
2312 i += 1;
2313 }
2314 text.trim_end()
2316}
2317
2318pub(crate) fn excerpt_pub(text: &str) -> String {
2321 excerpt(text)
2322}
2323
2324fn excerpt(text: &str) -> String {
2325 let value = one_line(text);
2326 value.chars().take(256).collect()
2327}
2328
2329fn one_line(text: &str) -> String {
2330 text.split_whitespace().collect::<Vec<_>>().join(" ")
2331}
2332
2333const FILE_EXPAND_CAP_BYTES: usize = 2048;
2340
2341pub fn resolve_capsule(
2353 conn: &Connection,
2354 repo_root: &std::path::Path,
2355 handle: &str,
2356) -> kimetsu_core::KimetsuResult<String> {
2357 if let Some(memory_id) = handle.strip_prefix("memory:") {
2358 let mut stmt = conn.prepare_cached(
2360 "SELECT text FROM memories WHERE memory_id = ? AND invalidated_at IS NULL",
2361 )?;
2362 let text: Option<String> = stmt
2363 .query_row(rusqlite::params![memory_id], |row| row.get(0))
2364 .optional()?;
2365 match text {
2366 Some(t) => Ok(t),
2367 None => {
2368 Err(format!("expand_capsule: no active memory found for handle `{handle}`").into())
2369 }
2370 }
2371 } else if let Some(rel_path) = handle.strip_prefix("file:") {
2372 let path = std::path::Path::new(rel_path);
2377 if path.is_absolute() {
2378 return Err(format!(
2379 "expand_capsule: `{handle}` is an absolute path — only repo-relative paths are supported"
2380 )
2381 .into());
2382 }
2383 for component in path.components() {
2384 match component {
2385 std::path::Component::ParentDir => {
2386 return Err(format!(
2387 "expand_capsule: `{handle}` contains `..` traversal — rejected"
2388 )
2389 .into());
2390 }
2391 std::path::Component::RootDir | std::path::Component::Prefix(_) => {
2392 return Err(format!(
2393 "expand_capsule: `{handle}` is an absolute path — only repo-relative paths are supported"
2394 )
2395 .into());
2396 }
2397 _ => {}
2398 }
2399 }
2400 let full_path = repo_root.join(path);
2401 let bytes = std::fs::read(&full_path)
2402 .map_err(|e| format!("expand_capsule: could not read `{rel_path}`: {e}"))?;
2403 let bounded = if bytes.len() > FILE_EXPAND_CAP_BYTES {
2405 let mut end = FILE_EXPAND_CAP_BYTES;
2406 while end > 0 && (bytes[end] & 0xC0) == 0x80 {
2408 end -= 1;
2409 }
2410 let s = String::from_utf8_lossy(&bytes[..end]);
2411 format!(
2412 "{s}\n[... truncated at {FILE_EXPAND_CAP_BYTES} bytes; call expand_capsule again with a line range if needed]"
2413 )
2414 } else {
2415 String::from_utf8_lossy(&bytes).into_owned()
2416 };
2417 Ok(bounded)
2418 } else if handle.starts_with("run:") {
2419 Err(format!(
2420 "expand_capsule: `run:` handle expansion is not yet supported (handle: `{handle}`)"
2421 )
2422 .into())
2423 } else {
2424 Err(format!(
2425 "expand_capsule: unrecognised handle format `{handle}`; \
2426 expected `memory:<id>`, `file:<path>`, or `run:<id>`"
2427 )
2428 .into())
2429 }
2430}
2431
2432pub fn rerank_capsules(
2441 query: &str,
2442 capsules: Vec<ContextCapsule>,
2443 reranker: &dyn crate::embeddings::Reranker,
2444 floor: f32,
2445 cap: usize,
2446) -> Vec<ContextCapsule> {
2447 if capsules.is_empty() {
2448 return capsules;
2449 }
2450
2451 let docs: Vec<&str> = capsules.iter().map(|c| c.summary.as_str()).collect();
2457 let scores = match reranker.rerank(query, &docs) {
2458 Ok(s) if s.len() == docs.len() => s,
2463 _ => {
2464 let mut out = capsules;
2466 if cap > 0 && out.len() > cap {
2467 out.truncate(cap);
2468 }
2469 return out;
2470 }
2471 };
2472
2473 let mut ranked: Vec<ContextCapsule> = capsules
2474 .into_iter()
2475 .zip(scores)
2476 .map(|(mut c, s)| {
2477 c.score = s;
2478 c
2479 })
2480 .collect();
2481
2482 ranked.sort_by(|a, b| {
2483 b.score
2484 .partial_cmp(&a.score)
2485 .unwrap_or(std::cmp::Ordering::Equal)
2486 });
2487
2488 ranked.retain(|c| c.score >= floor);
2489
2490 if cap > 0 && ranked.len() > cap {
2491 ranked.truncate(cap);
2492 }
2493
2494 ranked
2495}
2496
2497#[cfg(test)]
2498mod tests {
2499 use super::*;
2500
2501 fn capsule(kind: &str, summary: &str) -> ContextCapsule {
2502 ContextCapsule {
2503 id: "c".into(),
2504 kind: kind.into(),
2505 summary: summary.into(),
2506 token_estimate: 1,
2507 expansion_handle: "memory:x".into(),
2508 provenance: vec![],
2509 confidence: 1.0,
2510 freshness: 1.0,
2511 relevance: 1.0,
2512 scope_weight: 1.0,
2513 score: 1.0,
2514 }
2515 }
2516
2517 fn make_test_dir(tag: &str) -> std::path::PathBuf {
2520 use std::time::{SystemTime, UNIX_EPOCH};
2521 let ts = SystemTime::now()
2522 .duration_since(UNIX_EPOCH)
2523 .map(|d| d.subsec_nanos())
2524 .unwrap_or(0);
2525 let dir = std::env::temp_dir().join(format!("kbrain_test_{tag}_{ts}"));
2526 std::fs::create_dir_all(&dir).expect("create test dir");
2527 dir
2528 }
2529
2530 #[test]
2531 fn capsule_matches_kind_reads_memory_summary_prefix() {
2532 let mem = capsule("memory", "project:failure_pattern - linker not found");
2534 assert!(capsule_matches_kind(&mem, "failure_pattern"));
2535 assert!(!capsule_matches_kind(&mem, "command"));
2536 let repo = capsule("repo_file", "src/lib.rs:command - run build");
2538 assert!(capsule_matches_kind(&repo, "repo_file"));
2539 assert!(!capsule_matches_kind(&repo, "command"));
2540 }
2541
2542 #[test]
2545 fn usefulness_multiplier_neutral_at_zero_uses() {
2546 assert!((usefulness_multiplier(0.0, 0) - 1.0).abs() < f32::EPSILON);
2548 assert!((usefulness_multiplier(5.0, 0) - 1.0).abs() < f32::EPSILON);
2549 assert!((usefulness_multiplier(-5.0, 0) - 1.0).abs() < f32::EPSILON);
2550 }
2551
2552 #[test]
2556 fn usefulness_multiplier_blends_smoothly_in_transition() {
2557 let one_use = usefulness_multiplier(1.0, 1);
2560 assert!((one_use - 1.166_666_6).abs() < 1e-4, "got {one_use}");
2561 let two_uses = usefulness_multiplier(2.0, 2);
2564 assert!((two_uses - 1.333_333_4).abs() < 1e-4, "got {two_uses}");
2565 let two_uses_bad = usefulness_multiplier(-2.0, 2);
2567 assert!(
2569 (two_uses_bad - 0.666_666_7).abs() < 1e-4,
2570 "got {two_uses_bad}"
2571 );
2572 }
2573
2574 #[test]
2578 fn usefulness_multiplier_maps_ratio_onto_envelope() {
2579 assert!((usefulness_multiplier(5.0, 5) - 1.5).abs() < f32::EPSILON);
2581 assert!((usefulness_multiplier(-5.0, 5) - 0.5).abs() < f32::EPSILON);
2583 let mid = usefulness_multiplier(0.0, 6);
2585 assert!((mid - 1.0).abs() < f32::EPSILON, "got {mid}");
2586 let high = usefulness_multiplier(2.0, 4);
2588 assert!((high - 1.25).abs() < f32::EPSILON, "got {high}");
2589 let low = usefulness_multiplier(-2.0, 4);
2591 assert!((low - 0.75).abs() < f32::EPSILON, "got {low}");
2592 }
2593
2594 #[test]
2598 fn usefulness_multiplier_clamps_to_envelope() {
2599 assert!((usefulness_multiplier(100.0, 5) - 1.5).abs() < f32::EPSILON);
2601 assert!((usefulness_multiplier(-100.0, 5) - 0.5).abs() < f32::EPSILON);
2603 }
2604
2605 #[test]
2612 fn boost_gain_is_capped_so_cited_junk_cannot_beat_relevant_uncited() {
2613 let junk = apply_usefulness_boost(0.39, 1.5);
2614 let true_match = apply_usefulness_boost(0.53, 1.0);
2615 assert!(
2616 junk < true_match,
2617 "capped boost must preserve relevance order: junk {junk} vs match {true_match}"
2618 );
2619 assert!(junk <= 0.39 + USEFULNESS_BOOST_CAP + f32::EPSILON);
2621 }
2622
2623 #[test]
2627 fn boost_still_reorders_within_a_relevance_band() {
2628 let proven = apply_usefulness_boost(0.85, 1.5);
2629 let neutral = apply_usefulness_boost(0.90, 1.0);
2630 assert!(
2631 proven > neutral,
2632 "capped boost must still reorder near-equals: proven {proven} vs neutral {neutral}"
2633 );
2634 }
2635
2636 #[test]
2640 fn penalty_side_remains_multiplicative() {
2641 let penalized = apply_usefulness_boost(0.8, 0.5);
2642 assert!((penalized - 0.4).abs() < 1e-6);
2643 }
2644
2645 #[test]
2648 fn query_tokens_expands_build_class() {
2649 let toks = query_tokens("Build the project from source");
2650 assert!(toks.iter().any(|t| t == "build"));
2651 assert!(toks.iter().any(|t| t == "shell_background"));
2653 assert!(toks.iter().any(|t| t == "long_running"));
2654 }
2655
2656 #[test]
2657 fn query_tokens_expands_edit_class() {
2658 let toks = query_tokens("Modify the config to fix the bug");
2659 assert!(toks.iter().any(|t| t == "edit_file"));
2660 assert!(toks.iter().any(|t| t == "apply_patch"));
2661 }
2662
2663 #[test]
2664 fn query_tokens_expands_search_class() {
2665 let toks = query_tokens("Find all references to the symbol");
2666 assert!(toks.iter().any(|t| t == "glob"));
2667 assert!(toks.iter().any(|t| t == "search_files"));
2668 }
2669
2670 #[test]
2671 fn query_tokens_no_expansion_on_unrelated_query() {
2672 let toks = query_tokens("hello world testing nothing");
2673 assert!(toks.iter().any(|t| t == "hello"));
2675 assert!(toks.iter().any(|t| t == "world"));
2677 }
2678
2679 #[test]
2682 fn jaccard_is_zero_for_disjoint_sets() {
2683 let a: std::collections::HashSet<String> =
2684 ["foo", "bar"].iter().map(|s| s.to_string()).collect();
2685 let b: std::collections::HashSet<String> =
2686 ["baz", "qux"].iter().map(|s| s.to_string()).collect();
2687 assert!((jaccard(&a, &b) - 0.0).abs() < f32::EPSILON);
2688 }
2689
2690 #[test]
2691 fn jaccard_is_one_for_identical_sets() {
2692 let a: std::collections::HashSet<String> =
2693 ["foo", "bar"].iter().map(|s| s.to_string()).collect();
2694 let b = a.clone();
2695 assert!((jaccard(&a, &b) - 1.0).abs() < f32::EPSILON);
2696 }
2697
2698 #[test]
2699 fn jaccard_partial_overlap() {
2700 let a: std::collections::HashSet<String> = ["foo", "bar", "baz"]
2701 .iter()
2702 .map(|s| s.to_string())
2703 .collect();
2704 let b: std::collections::HashSet<String> =
2705 ["bar", "qux"].iter().map(|s| s.to_string()).collect();
2706 assert!((jaccard(&a, &b) - 0.25).abs() < f32::EPSILON);
2708 }
2709
2710 #[test]
2711 fn summary_token_set_lowercases_and_filters_short() {
2712 let set = summary_token_set("Build the Foo-bar project");
2713 assert!(set.contains("build"));
2714 assert!(set.contains("foo"));
2715 assert!(set.contains("bar"));
2716 assert!(set.contains("project"));
2717 assert!(set.contains("the"));
2719 }
2720
2721 fn insert_memory_with_embedding(
2727 conn: &rusqlite::Connection,
2728 memory_id: &str,
2729 text: &str,
2730 embedder: &dyn embeddings::Embedder,
2731 ) {
2732 let normalized = kimetsu_core::memory::normalize_memory_text(text);
2733 conn.execute(
2734 "
2735 INSERT INTO memories (
2736 memory_id, scope, kind, text, normalized_text, confidence,
2737 source_event_id, provenance_snapshot_json, created_at,
2738 use_count, usefulness_score, embedding, embedding_model
2739 )
2740 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
2741 '2026-05-01T00:00:00Z', 0, 0.0, ?4, ?5)
2742 ",
2743 rusqlite::params![
2744 memory_id,
2745 text,
2746 normalized,
2747 embeddings::encode_embedding(&embedder.embed(text).expect("embed test row")),
2748 embedder.model_id(),
2749 ],
2750 )
2751 .expect("insert memory");
2752 conn.execute(
2753 "INSERT INTO memories_fts (memory_id, text, kind, scope) VALUES (?1, ?2, 'fact', 'global_user')",
2754 rusqlite::params![memory_id, text],
2755 )
2756 .expect("insert fts row");
2757 }
2758
2759 #[test]
2769 fn hybrid_retrieval_uses_cosine_score_to_rerank() {
2770 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
2771 crate::schema::initialize(&conn).expect("init schema");
2772 let stub = embeddings::StubEmbedder::new();
2773
2774 insert_memory_with_embedding(&conn, "m_rg", "use ripgrep for code search", &stub);
2775 insert_memory_with_embedding(
2776 &conn,
2777 "m_unrelated",
2778 "cookie recipe with chocolate chips",
2779 &stub,
2780 );
2781
2782 let weights = kimetsu_core::config::BrokerWeights::default();
2785 let bundle = retrieve_context_with_embedder(
2786 &conn,
2787 "/fake-repo",
2788 &weights,
2789 ContextRequest {
2790 stage: "localization".to_string(),
2791 query: "ripgrep search".to_string(),
2792 budget_tokens: 4000,
2793 ..Default::default()
2794 },
2795 &[],
2796 &stub,
2797 )
2798 .expect("retrieve");
2799
2800 let memory_handles: Vec<_> = bundle
2801 .capsules
2802 .iter()
2803 .filter(|c| c.expansion_handle.starts_with("memory:"))
2804 .collect();
2805 assert!(
2806 !memory_handles.is_empty(),
2807 "at least one memory should surface"
2808 );
2809 assert_eq!(
2811 memory_handles[0].expansion_handle,
2812 "memory:m_rg",
2813 "ripgrep memory should outrank the cookie recipe; ranked: {:?}",
2814 memory_handles
2815 .iter()
2816 .map(|c| &c.expansion_handle)
2817 .collect::<Vec<_>>()
2818 );
2819 }
2820
2821 #[test]
2828 fn hybrid_retrieval_skips_cosine_on_model_id_mismatch() {
2829 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
2830 crate::schema::initialize(&conn).expect("init schema");
2831 let stub = embeddings::StubEmbedder::new();
2832 insert_memory_with_embedding(&conn, "m_xref", "use ripgrep for code search", &stub);
2833
2834 conn.execute(
2839 "UPDATE memories SET embedding_model = 'bge-small-en-v1.5' WHERE memory_id = 'm_xref'",
2840 [],
2841 )
2842 .expect("force model_id mismatch");
2843
2844 let weights = kimetsu_core::config::BrokerWeights::default();
2849 let bundle = retrieve_context_with_embedder(
2850 &conn,
2851 "/fake-repo",
2852 &weights,
2853 ContextRequest {
2854 stage: "localization".to_string(),
2855 query: "ripgrep search".to_string(),
2856 budget_tokens: 4000,
2857 ..Default::default()
2858 },
2859 &[],
2860 &stub,
2861 )
2862 .expect("retrieve");
2863
2864 assert!(
2865 bundle
2866 .capsules
2867 .iter()
2868 .any(|c| c.expansion_handle == "memory:m_xref"),
2869 "cross-model row should still match lexically (cosine skipped, FTS works)"
2870 );
2871 }
2872
2873 #[test]
2880 fn usefulness_decay_disabled_when_half_life_is_zero_or_negative() {
2881 let ancient = "2021-01-01T00:00:00Z";
2883 assert!((usefulness_decay(Some(ancient), ancient, 0.0) - 1.0).abs() < f32::EPSILON);
2884 assert!((usefulness_decay(Some(ancient), ancient, -1.0) - 1.0).abs() < f32::EPSILON);
2885 }
2886
2887 #[test]
2891 fn usefulness_decay_returns_one_on_unparseable_timestamps() {
2892 assert!(
2893 (usefulness_decay(Some("not-a-date"), "also-not", 30.0) - 1.0).abs() < f32::EPSILON
2894 );
2895 }
2896
2897 #[test]
2900 fn usefulness_decay_full_at_zero_age() {
2901 let future = "2099-01-01T00:00:00Z";
2903 let d = usefulness_decay(Some(future), future, 30.0);
2904 assert!((d - 1.0).abs() < f32::EPSILON, "got {d}");
2905 }
2906
2907 #[test]
2912 fn usefulness_decay_follows_half_life_curve() {
2913 let half_life = 10.0_f32;
2914 let now = OffsetDateTime::now_utc();
2915 let fmt = &time::format_description::well_known::Rfc3339;
2916
2917 let one_half_life_ago = (now - time::Duration::seconds((half_life * 86_400.0) as i64))
2919 .format(fmt)
2920 .expect("format");
2921 let d1 = usefulness_decay(Some(&one_half_life_ago), &one_half_life_ago, half_life);
2922 assert!(
2923 (d1 - 0.5).abs() < 0.01,
2924 "expected ~0.5 at one half-life, got {d1}"
2925 );
2926
2927 let two_half_lives_ago = (now
2929 - time::Duration::seconds((2.0 * half_life * 86_400.0) as i64))
2930 .format(fmt)
2931 .expect("format");
2932 let d2 = usefulness_decay(Some(&two_half_lives_ago), &two_half_lives_ago, half_life);
2933 assert!(
2934 (d2 - 0.25).abs() < 0.01,
2935 "expected ~0.25 at two half-lives, got {d2}"
2936 );
2937 }
2938
2939 #[test]
2943 fn usefulness_decay_falls_back_to_created_at_when_last_useful_is_none() {
2944 let now = OffsetDateTime::now_utc();
2945 let fmt = &time::format_description::well_known::Rfc3339;
2946 let one_day_ago = (now - time::Duration::seconds(86_400))
2947 .format(fmt)
2948 .expect("format");
2949 let d = usefulness_decay(None, &one_day_ago, 30.0);
2950 assert!(
2952 (d - 0.977).abs() < 0.01,
2953 "expected ~0.977 for 1-day-old created_at under 30d half-life, got {d}"
2954 );
2955 }
2956
2957 #[test]
2962 fn aged_cited_memory_ranks_below_recently_cited_memory() {
2963 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
2964 crate::schema::initialize(&conn).expect("init schema");
2965
2966 let now = OffsetDateTime::now_utc();
2967 let fmt = &time::format_description::well_known::Rfc3339;
2968 let one_day_ago = (now - time::Duration::seconds(86_400))
2969 .format(fmt)
2970 .expect("format");
2971 let one_year_ago = (now - time::Duration::seconds(365 * 86_400))
2972 .format(fmt)
2973 .expect("format");
2974
2975 for (mid, last_useful) in [("m_recent", &one_day_ago), ("m_aged", &one_year_ago)] {
2979 let text = "use ripgrep for code search";
2980 let normalized = kimetsu_core::memory::normalize_memory_text(text);
2981 conn.execute(
2982 "
2983 INSERT INTO memories (
2984 memory_id, scope, kind, text, normalized_text, confidence,
2985 source_event_id, provenance_snapshot_json, created_at,
2986 use_count, usefulness_score, last_useful_at
2987 )
2988 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
2989 '2024-01-01T00:00:00Z', 5, 5.0, ?4)
2990 ",
2991 rusqlite::params![mid, text, normalized, last_useful],
2992 )
2993 .expect("insert memory");
2994 conn.execute(
2995 "INSERT INTO memories_fts (memory_id, text, kind, scope)
2996 VALUES (?1, ?2, 'fact', 'global_user')",
2997 rusqlite::params![mid, text],
2998 )
2999 .expect("insert fts");
3000 }
3001
3002 let weights = kimetsu_core::config::BrokerWeights::default();
3004 let bundle = retrieve_context_with_embedder(
3005 &conn,
3006 "/fake-repo",
3007 &weights,
3008 ContextRequest {
3009 stage: "localization".to_string(),
3010 query: "ripgrep search".to_string(),
3011 budget_tokens: 4000,
3012 ..Default::default()
3013 },
3014 &[],
3015 &embeddings::NoopEmbedder,
3016 )
3017 .expect("retrieve");
3018
3019 let mem_order: Vec<&str> = bundle
3020 .capsules
3021 .iter()
3022 .filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
3023 .collect();
3024 assert_eq!(
3025 mem_order.first().copied(),
3026 Some("m_recent"),
3027 "recently-cited memory must rank first under decay; got order {mem_order:?}"
3028 );
3029 }
3030
3031 #[test]
3036 fn aged_cited_memory_does_not_decay_when_half_life_is_zero() {
3037 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3038 crate::schema::initialize(&conn).expect("init schema");
3039
3040 let now = OffsetDateTime::now_utc();
3041 let fmt = &time::format_description::well_known::Rfc3339;
3042 let one_day_ago = (now - time::Duration::seconds(86_400))
3043 .format(fmt)
3044 .expect("format");
3045 let one_year_ago = (now - time::Duration::seconds(365 * 86_400))
3046 .format(fmt)
3047 .expect("format");
3048
3049 for (mid, last_useful) in [("m_recent", &one_day_ago), ("m_aged", &one_year_ago)] {
3050 let text = "use ripgrep for code search";
3051 let normalized = kimetsu_core::memory::normalize_memory_text(text);
3052 conn.execute(
3053 "
3054 INSERT INTO memories (
3055 memory_id, scope, kind, text, normalized_text, confidence,
3056 source_event_id, provenance_snapshot_json, created_at,
3057 use_count, usefulness_score, last_useful_at
3058 )
3059 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
3060 '2024-01-01T00:00:00Z', 5, 5.0, ?4)
3061 ",
3062 rusqlite::params![mid, text, normalized, last_useful],
3063 )
3064 .expect("insert memory");
3065 conn.execute(
3066 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3067 VALUES (?1, ?2, 'fact', 'global_user')",
3068 rusqlite::params![mid, text],
3069 )
3070 .expect("insert fts");
3071 }
3072
3073 let weights = kimetsu_core::config::BrokerWeights {
3075 decay_half_life_days: 0.0,
3076 ..Default::default()
3077 };
3078
3079 let bundle = retrieve_context_with_embedder(
3080 &conn,
3081 "/fake-repo",
3082 &weights,
3083 ContextRequest {
3084 stage: "localization".to_string(),
3085 query: "ripgrep search".to_string(),
3086 budget_tokens: 4000,
3087 ..Default::default()
3088 },
3089 &[],
3090 &embeddings::NoopEmbedder,
3091 )
3092 .expect("retrieve");
3093
3094 let scores: Vec<(String, f32)> = bundle
3100 .capsules
3101 .iter()
3102 .filter_map(|c| {
3103 c.expansion_handle
3104 .strip_prefix("memory:")
3105 .map(|id| (id.to_string(), c.score))
3106 })
3107 .collect();
3108 assert_eq!(scores.len(), 2, "both memories should surface");
3109 let recent_score = scores
3110 .iter()
3111 .find(|(id, _)| id == "m_recent")
3112 .map(|(_, s)| *s)
3113 .expect("m_recent present");
3114 let aged_score = scores
3115 .iter()
3116 .find(|(id, _)| id == "m_aged")
3117 .map(|(_, s)| *s)
3118 .expect("m_aged present");
3119 assert!(
3121 (recent_score - aged_score).abs() < 1e-4,
3122 "with decay disabled the two memories should tie on score: recent={recent_score} aged={aged_score}"
3123 );
3124 }
3125
3126 #[test]
3131 fn hybrid_retrieval_with_noop_embedder_is_lexical_only() {
3132 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3133 crate::schema::initialize(&conn).expect("init schema");
3134 let stub = embeddings::StubEmbedder::new();
3135 insert_memory_with_embedding(&conn, "m_a", "use ripgrep", &stub);
3137 insert_memory_with_embedding(&conn, "m_b", "use ripgrep too", &stub);
3138
3139 let weights = kimetsu_core::config::BrokerWeights::default();
3142 let bundle = retrieve_context_with_embedder(
3143 &conn,
3144 "/fake-repo",
3145 &weights,
3146 ContextRequest {
3147 stage: "localization".to_string(),
3148 query: "ripgrep".to_string(),
3149 budget_tokens: 4000,
3150 ..Default::default()
3151 },
3152 &[],
3153 &embeddings::NoopEmbedder,
3154 )
3155 .expect("retrieve");
3156
3157 let count = bundle
3158 .capsules
3159 .iter()
3160 .filter(|c| c.expansion_handle.starts_with("memory:"))
3161 .count();
3162 assert_eq!(count, 2, "both memories should surface via FTS");
3163 }
3164
3165 #[cfg(feature = "embeddings")]
3192 #[test]
3193 fn ann_finds_semantic_match_fts_misses() {
3194 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3195 crate::schema::initialize(&conn).expect("init schema");
3196
3197 struct OracleEmbedder;
3200 impl embeddings::Embedder for OracleEmbedder {
3201 fn embed(&self, _text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
3202 Ok(vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
3204 }
3205 fn model_id(&self) -> &str {
3206 "oracle-d8"
3207 }
3208 fn dim(&self) -> usize {
3209 8
3210 }
3211 }
3212
3213 let model_id = "oracle-d8";
3214
3215 let sem_vec = embeddings::encode_embedding(&[1.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
3218 let sem_text = "cookie recipe chocolate";
3219 let sem_norm = kimetsu_core::memory::normalize_memory_text(sem_text);
3220 conn.execute(
3221 "INSERT INTO memories (
3222 memory_id, scope, kind, text, normalized_text, confidence,
3223 source_event_id, provenance_snapshot_json, created_at,
3224 use_count, usefulness_score, embedding, embedding_model
3225 )
3226 VALUES ('m_semantic', 'global_user', 'fact', ?1, ?2, 1.0, NULL, '{}',
3227 '2026-01-01T00:00:00Z', 0, 0.0, ?3, ?4)",
3228 rusqlite::params![sem_text, sem_norm, sem_vec, model_id],
3229 )
3230 .expect("insert m_semantic");
3231 conn.execute(
3232 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3233 VALUES ('m_semantic', ?1, 'fact', 'global_user')",
3234 rusqlite::params![sem_text],
3235 )
3236 .expect("insert m_semantic fts");
3237
3238 let decoy_vec = embeddings::encode_embedding(&[0.0f32, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
3240 let decoy_text = "git rebase squash commits";
3241 let decoy_norm = kimetsu_core::memory::normalize_memory_text(decoy_text);
3242 conn.execute(
3243 "INSERT INTO memories (
3244 memory_id, scope, kind, text, normalized_text, confidence,
3245 source_event_id, provenance_snapshot_json, created_at,
3246 use_count, usefulness_score, embedding, embedding_model
3247 )
3248 VALUES ('m_decoy', 'global_user', 'fact', ?1, ?2, 1.0, NULL, '{}',
3249 '2026-01-01T00:00:00Z', 0, 0.0, ?3, ?4)",
3250 rusqlite::params![decoy_text, decoy_norm, decoy_vec, model_id],
3251 )
3252 .expect("insert m_decoy");
3253 conn.execute(
3254 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3255 VALUES ('m_decoy', ?1, 'fact', 'global_user')",
3256 rusqlite::params![decoy_text],
3257 )
3258 .expect("insert m_decoy fts");
3259
3260 let fts_hits: i64 = conn
3262 .query_row(
3263 "SELECT COUNT(*) FROM memories_fts \
3264 WHERE memories_fts MATCH 'phosphorescent bioluminescent'",
3265 [],
3266 |r| r.get(0),
3267 )
3268 .unwrap_or(0);
3269 assert_eq!(
3270 fts_hits, 0,
3271 "sanity: query tokens must not appear in any memory text"
3272 );
3273
3274 let weights = kimetsu_core::config::BrokerWeights::default();
3278 let bundle = retrieve_context_with_embedder(
3279 &conn,
3280 "/fake-repo",
3281 &weights,
3282 ContextRequest {
3283 stage: "localization".to_string(),
3284 query: "phosphorescent bioluminescent organism".to_string(),
3285 budget_tokens: 4000,
3286 ..Default::default()
3287 },
3288 &[],
3289 &OracleEmbedder,
3290 )
3291 .expect("retrieve");
3292
3293 let handles: Vec<&str> = bundle
3294 .capsules
3295 .iter()
3296 .filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
3297 .collect();
3298
3299 assert!(
3300 handles.contains(&"m_semantic"),
3301 "ANN must surface m_semantic (cosine=1 with oracle query) even though \
3302 FTS found nothing; got handles: {handles:?}"
3303 );
3304 }
3305
3306 #[cfg(feature = "embeddings")]
3308 #[test]
3309 fn dedup_memory_matched_by_fts_and_ann_appears_once() {
3310 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3311 crate::schema::initialize(&conn).expect("init schema");
3312
3313 let stub = embeddings::StubEmbedder::new();
3314
3315 insert_memory_with_embedding(&conn, "m_both", "use ripgrep for fast search", &stub);
3319
3320 let weights = kimetsu_core::config::BrokerWeights::default();
3321 let bundle = retrieve_context_with_embedder(
3322 &conn,
3323 "/fake-repo",
3324 &weights,
3325 ContextRequest {
3326 stage: "localization".to_string(),
3327 query: "ripgrep".to_string(),
3328 budget_tokens: 4000,
3329 ..Default::default()
3330 },
3331 &[],
3332 &stub,
3333 )
3334 .expect("retrieve");
3335
3336 let count = bundle
3337 .capsules
3338 .iter()
3339 .filter(|c| c.expansion_handle == "memory:m_both")
3340 .count();
3341 assert_eq!(
3342 count,
3343 1,
3344 "m_both (matched by both FTS and ANN) must appear exactly once; \
3345 bundle: {:?}",
3346 bundle
3347 .capsules
3348 .iter()
3349 .map(|c| &c.expansion_handle)
3350 .collect::<Vec<_>>()
3351 );
3352 }
3353
3354 #[cfg(feature = "embeddings")]
3378 #[test]
3379 fn embedding_mmr_collapses_paraphrases_but_jaccard_does_not() {
3380 struct OracleEmbedder;
3383 impl embeddings::Embedder for OracleEmbedder {
3384 fn embed(&self, _text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
3385 let mut v = vec![0.0f32; 8];
3386 v[0] = 1.0;
3387 Ok(v)
3388 }
3389 fn model_id(&self) -> &str {
3390 "oracle-d8"
3391 }
3392 fn dim(&self) -> usize {
3393 8
3394 }
3395 }
3396
3397 let oracle = OracleEmbedder;
3400 let weights = kimetsu_core::config::BrokerWeights::default();
3401
3402 let m_rg1_text = "prefer ripgrep for searching source code";
3405 let m_rg2_text = "rg is the fastest way to locate patterns";
3406
3407 let conn = rusqlite::Connection::open_in_memory().expect("in-memory");
3413 crate::schema::initialize(&conn).expect("init schema");
3414 insert_memory_with_embedding(&conn, "m_rg1", m_rg1_text, &oracle);
3415 insert_memory_with_embedding(&conn, "m_rg2", m_rg2_text, &oracle);
3416
3417 let bundle_embedding = retrieve_context_with_embedder(
3418 &conn,
3419 "/fake-repo",
3420 &weights,
3421 ContextRequest {
3422 stage: "localization".to_string(),
3423 query: "search source patterns".to_string(),
3425 budget_tokens: 20_000,
3426 max_capsules: 1, ..Default::default()
3428 },
3429 &[],
3430 &oracle,
3431 )
3432 .expect("retrieve with oracle embedder");
3433
3434 let emb_in_capsules = bundle_embedding
3437 .capsules
3438 .iter()
3439 .filter(|c| {
3440 c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
3441 })
3442 .count();
3443 assert_eq!(
3444 emb_in_capsules,
3445 1,
3446 "embedding-MMR must collapse cosine=1.0 paraphrases: with max_capsules=1 \
3447 only ONE should be included; capsule handles: {:?}; excluded: {:?}",
3448 bundle_embedding
3449 .capsules
3450 .iter()
3451 .map(|c| &c.expansion_handle)
3452 .collect::<Vec<_>>(),
3453 bundle_embedding
3454 .excluded
3455 .iter()
3456 .map(|c| &c.expansion_handle)
3457 .collect::<Vec<_>>()
3458 );
3459
3460 let emb_in_excluded = bundle_embedding
3462 .excluded
3463 .iter()
3464 .filter(|c| {
3465 c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
3466 })
3467 .count();
3468 assert_eq!(
3469 emb_in_excluded,
3470 1,
3471 "the second near-duplicate must be in excluded under embedding-MMR; \
3472 excluded handles: {:?}",
3473 bundle_embedding
3474 .excluded
3475 .iter()
3476 .map(|c| &c.expansion_handle)
3477 .collect::<Vec<_>>()
3478 );
3479
3480 let conn2 = rusqlite::Connection::open_in_memory().expect("in-memory 2");
3485 crate::schema::initialize(&conn2).expect("init schema 2");
3486 insert_memory_with_embedding(&conn2, "m_rg1", m_rg1_text, &oracle);
3487 insert_memory_with_embedding(&conn2, "m_rg2", m_rg2_text, &oracle);
3488
3489 let bundle_lean = retrieve_context_with_embedder(
3490 &conn2,
3491 "/fake-repo",
3492 &weights,
3493 ContextRequest {
3494 stage: "localization".to_string(),
3495 query: "search source patterns".to_string(),
3496 budget_tokens: 20_000,
3497 max_capsules: 2, ..Default::default()
3499 },
3500 &[],
3501 &embeddings::NoopEmbedder,
3502 )
3503 .expect("retrieve with NoopEmbedder");
3504
3505 let lean_in_capsules = bundle_lean
3506 .capsules
3507 .iter()
3508 .filter(|c| {
3509 c.expansion_handle == "memory:m_rg1" || c.expansion_handle == "memory:m_rg2"
3510 })
3511 .count();
3512 assert_eq!(
3513 lean_in_capsules,
3514 2,
3515 "Jaccard-only path must NOT collapse the two paraphrases (different words, \
3516 low token overlap → both survive MMR with max_capsules=2); capsule handles: {:?}",
3517 bundle_lean
3518 .capsules
3519 .iter()
3520 .map(|c| &c.expansion_handle)
3521 .collect::<Vec<_>>()
3522 );
3523 }
3524
3525 #[test]
3528 fn content_tokens_strips_stopwords_keeps_topical_words() {
3529 let got = content_tokens("Tell me about kimetsu, what's the idea of the repo");
3530 assert_eq!(got, vec!["kimetsu", "idea", "repo"]);
3533 }
3534
3535 #[test]
3536 fn light_stem_strips_one_inflection_suffix() {
3537 assert_eq!(light_stem("benchmarked"), "benchmark");
3538 assert_eq!(light_stem("benchmarking"), "benchmark");
3539 assert_eq!(light_stem("repos"), "repo");
3540 assert_eq!(light_stem("does"), "does");
3542 assert_eq!(light_stem("toml"), "toml");
3543 }
3544
3545 #[test]
3552 fn stemmed_query_matches_inflected_corpus_through_floor() {
3553 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3554 crate::schema::initialize(&conn).expect("init schema");
3555 let insert = |id: &str, text: &str| {
3556 let norm = kimetsu_core::memory::normalize_memory_text(text);
3557 conn.execute(
3558 "INSERT INTO memories (
3559 memory_id, scope, kind, text, normalized_text, confidence,
3560 source_event_id, provenance_snapshot_json, created_at,
3561 use_count, usefulness_score, embedding, embedding_model
3562 )
3563 VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
3564 '2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
3565 rusqlite::params![id, text, norm],
3566 )
3567 .expect("insert memory");
3568 conn.execute(
3569 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3570 VALUES (?1, ?2, 'fact', 'global_user')",
3571 rusqlite::params![id, text],
3572 )
3573 .expect("insert fts");
3574 };
3575 insert(
3576 "m_bench",
3577 "kimetsu benchmark runs go through the kbench binary and the Terminal-Bench driver",
3578 );
3579 insert(
3580 "m_doctor",
3581 "kimetsu doctor version-skew check parses process start times on Windows via CIM",
3582 );
3583 insert(
3584 "m_gc",
3585 "kimetsu runs auto-GC on run creation; keep the env guard at the trigger site",
3586 );
3587
3588 let bundle = retrieve_context_with_embedder(
3589 &conn,
3590 "/fake-repo",
3591 &kimetsu_core::config::BrokerWeights::default(),
3592 ContextRequest {
3593 stage: "localization".to_string(),
3594 query: "Can you find out how kimetsu is benchmarked?".to_string(),
3595 budget_tokens: 2000,
3596 max_capsules: 2,
3597 min_lexical_coverage: 0.5,
3598 ..Default::default()
3599 },
3600 &[],
3601 &embeddings::NoopEmbedder,
3602 )
3603 .expect("retrieve");
3604 let handles: Vec<_> = bundle
3605 .capsules
3606 .iter()
3607 .map(|c| c.expansion_handle.as_str())
3608 .collect();
3609 assert!(
3610 handles.contains(&"memory:m_bench"),
3611 "stemmed 'benchmarked' must surface the benchmark memory; got {handles:?}"
3612 );
3613 assert!(
3614 !handles.contains(&"memory:m_doctor") && !handles.contains(&"memory:m_gc"),
3615 "off-topic memories sharing only 'kimetsu' must stay below the floor; got {handles:?}"
3616 );
3617 }
3618
3619 #[test]
3620 fn weighted_coverage_ignores_zero_idf_tokens() {
3621 let content = vec![
3625 "kimetsu".to_string(),
3626 "idea".to_string(),
3627 "repo".to_string(),
3628 ];
3629 let mut idf = HashMap::new();
3630 idf.insert("kimetsu".to_string(), 0.0);
3631 idf.insert("idea".to_string(), 1.386);
3632 idf.insert("repo".to_string(), 0.693);
3633
3634 let cov = weighted_coverage(
3636 &content,
3637 &idf,
3638 "global:fact - the git repo and kimetsu brain",
3639 );
3640 assert!((cov - 0.333).abs() < 0.01, "got {cov}");
3641
3642 let cov_topical =
3644 weighted_coverage(&content, &idf, "global:fact - the core idea of kimetsu");
3645 assert!(cov_topical > 0.6, "got {cov_topical}");
3646 }
3647
3648 #[test]
3649 fn escape_like_neutralizes_wildcards() {
3650 assert_eq!(escape_like("a_b%c"), "a\\_b\\%c");
3651 assert_eq!(escape_like("plain"), "plain");
3652 }
3653
3654 #[test]
3668 fn lexical_floor_drops_offtopic_memories_sharing_project_name() {
3669 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3670 crate::schema::initialize(&conn).expect("init schema");
3671
3672 let insert = |id: &str, text: &str| {
3673 let norm = kimetsu_core::memory::normalize_memory_text(text);
3674 conn.execute(
3675 "INSERT INTO memories (
3676 memory_id, scope, kind, text, normalized_text, confidence,
3677 source_event_id, provenance_snapshot_json, created_at,
3678 use_count, usefulness_score, embedding, embedding_model
3679 )
3680 VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
3681 '2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
3682 rusqlite::params![id, text, norm],
3683 )
3684 .expect("insert memory");
3685 conn.execute(
3686 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3687 VALUES (?1, ?2, 'fact', 'global_user')",
3688 rusqlite::params![id, text],
3689 )
3690 .expect("insert fts");
3691 };
3692
3693 insert(
3696 "m1",
3697 "When implementing a setup command that calls init_project, tests must call \
3698 git_init_boundary before setup_cmd so ProjectPaths discover resolves to the temp \
3699 dir instead of climbing to the real parent git repo including the user brain at kimetsu",
3700 );
3701 insert(
3702 "m2",
3703 "A member crate with default embeddings silently turned embeddings on for the entire \
3704 cargo test workspace build graph because cargo unifies features; kimetsu-chat \
3705 retrieval tests failed",
3706 );
3707 insert(
3708 "m3",
3709 "In toml 0.9 use toml from_str to parse a TOML document into a Value not str parse; \
3710 implementing config get and set in kimetsu-cli",
3711 );
3712
3713 let query = "Tell me about kimetsu, what's the idea of the repo".to_string();
3714 let weights = kimetsu_core::config::BrokerWeights::default();
3715 let handles = |bundle: &ContextBundle| {
3716 bundle
3717 .capsules
3718 .iter()
3719 .map(|c| c.expansion_handle.clone())
3720 .collect::<Vec<_>>()
3721 };
3722
3723 let no_floor = retrieve_context_with_embedder(
3725 &conn,
3726 "/fake-repo",
3727 &weights,
3728 ContextRequest {
3729 stage: "localization".to_string(),
3730 query: query.clone(),
3731 budget_tokens: 2000,
3732 max_capsules: 8,
3733 min_lexical_coverage: 0.0,
3734 ..Default::default()
3735 },
3736 &[],
3737 &embeddings::NoopEmbedder,
3738 )
3739 .expect("retrieve without floor");
3740 let before = handles(&no_floor);
3741 assert!(
3742 before.contains(&"memory:m2".to_string()) && before.contains(&"memory:m3".to_string()),
3743 "sanity: without the floor the pure-project-name memories should surface; got {before:?}"
3744 );
3745
3746 let floored = retrieve_context_with_embedder(
3748 &conn,
3749 "/fake-repo",
3750 &weights,
3751 ContextRequest {
3752 stage: "localization".to_string(),
3753 query,
3754 budget_tokens: 2000,
3755 max_capsules: 8,
3756 min_lexical_coverage: 0.5,
3757 ..Default::default()
3758 },
3759 &[],
3760 &embeddings::NoopEmbedder,
3761 )
3762 .expect("retrieve with floor");
3763 let after = handles(&floored);
3764 assert!(
3765 !after.contains(&"memory:m2".to_string()) && !after.contains(&"memory:m3".to_string()),
3766 "the lexical floor must drop memories whose only match is the corpus-ubiquitous \
3767 project name; surviving: {after:?}"
3768 );
3769 }
3770
3771 #[test]
3774 fn lexical_floor_keeps_ontopic_memory() {
3775 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
3776 crate::schema::initialize(&conn).expect("init schema");
3777
3778 let insert = |id: &str, text: &str| {
3779 let norm = kimetsu_core::memory::normalize_memory_text(text);
3780 conn.execute(
3781 "INSERT INTO memories (
3782 memory_id, scope, kind, text, normalized_text, confidence,
3783 source_event_id, provenance_snapshot_json, created_at,
3784 use_count, usefulness_score, embedding, embedding_model
3785 )
3786 VALUES (?1, 'global_user', 'fact', ?2, ?3, 0.9, NULL, '{}',
3787 '2026-06-01T00:00:00Z', 0, 0.0, NULL, NULL)",
3788 rusqlite::params![id, text, norm],
3789 )
3790 .expect("insert memory");
3791 conn.execute(
3792 "INSERT INTO memories_fts (memory_id, text, kind, scope)
3793 VALUES (?1, ?2, 'fact', 'global_user')",
3794 rusqlite::params![id, text],
3795 )
3796 .expect("insert fts");
3797 };
3798
3799 insert(
3801 "d1",
3802 "The distiller runs at session end and harvests durable lessons from the transcript",
3803 );
3804 insert(
3805 "n1",
3806 "Unrelated note about git rebase and squashing commits",
3807 );
3808
3809 let bundle = retrieve_context_with_embedder(
3810 &conn,
3811 "/fake-repo",
3812 &kimetsu_core::config::BrokerWeights::default(),
3813 ContextRequest {
3814 stage: "localization".to_string(),
3815 query: "how does the distiller work".to_string(),
3816 budget_tokens: 2000,
3817 min_lexical_coverage: 0.5,
3818 ..Default::default()
3819 },
3820 &[],
3821 &embeddings::NoopEmbedder,
3822 )
3823 .expect("retrieve");
3824
3825 assert!(
3826 bundle
3827 .capsules
3828 .iter()
3829 .any(|c| c.expansion_handle == "memory:d1"),
3830 "on-topic memory covering the rare query word must survive the floor; got: {:?}",
3831 bundle
3832 .capsules
3833 .iter()
3834 .map(|c| &c.expansion_handle)
3835 .collect::<Vec<_>>()
3836 );
3837 }
3838
3839 #[cfg(feature = "embeddings")]
3848 #[test]
3849 fn min_semantic_score_floor_drops_off_topic_queries() {
3850 struct DirectionalEmbedder {
3861 marker: &'static str,
3863 }
3864 impl embeddings::Embedder for DirectionalEmbedder {
3865 fn embed(&self, text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
3866 let mut v = vec![0.0f32; 8];
3867 if text.contains(self.marker) {
3868 v[0] = 1.0;
3869 } else {
3870 v[1] = 1.0;
3871 }
3872 Ok(v)
3873 }
3874 fn model_id(&self) -> &str {
3875 "directional-d8"
3876 }
3877 fn dim(&self) -> usize {
3878 8
3879 }
3880 }
3881
3882 let emb = DirectionalEmbedder { marker: "TOPIC_A" };
3883
3884 let conn = rusqlite::Connection::open_in_memory().expect("in-memory");
3885 crate::schema::initialize(&conn).expect("init schema");
3886
3887 insert_memory_with_embedding(&conn, "m_b", "cookie recipe chocolate baking TOPIC_B", &emb);
3889
3890 let weights = kimetsu_core::config::BrokerWeights::default();
3891
3892 let bundle_off = retrieve_context_with_embedder(
3894 &conn,
3895 "/fake-repo",
3896 &weights,
3897 ContextRequest {
3898 stage: "localization".to_string(),
3899 query: "TOPIC_A unrelated phosphorescent".to_string(),
3901 budget_tokens: 4000,
3902 min_semantic_score: 0.1, ..Default::default()
3904 },
3905 &[],
3906 &emb,
3907 )
3908 .expect("retrieve off-topic");
3909
3910 assert!(
3911 bundle_off.capsules.is_empty(),
3912 "off-topic query (cosine=0 < floor=0.1) must produce zero capsules; \
3913 got: {:?}",
3914 bundle_off
3915 .capsules
3916 .iter()
3917 .map(|c| &c.expansion_handle)
3918 .collect::<Vec<_>>()
3919 );
3920
3921 let conn2 = rusqlite::Connection::open_in_memory().expect("in-memory 2");
3924 crate::schema::initialize(&conn2).expect("init schema 2");
3925 insert_memory_with_embedding(
3926 &conn2,
3927 "m_b2",
3928 "cookie recipe chocolate TOPIC_B baking"
3929 .to_string()
3930 .as_str(),
3931 &emb,
3932 );
3933
3934 let bundle_on = retrieve_context_with_embedder(
3935 &conn2,
3936 "/fake-repo",
3937 &weights,
3938 ContextRequest {
3939 stage: "localization".to_string(),
3940 query: "cookie chocolate TOPIC_B".to_string(),
3942 budget_tokens: 4000,
3943 min_semantic_score: 0.1,
3944 ..Default::default()
3945 },
3946 &[],
3947 &emb,
3948 )
3949 .expect("retrieve on-topic");
3950
3951 assert!(
3952 bundle_on
3953 .capsules
3954 .iter()
3955 .any(|c| c.expansion_handle == "memory:m_b2"),
3956 "on-topic query (cosine=1.0 ≥ floor) must surface m_b2; \
3957 got capsules: {:?}",
3958 bundle_on
3959 .capsules
3960 .iter()
3961 .map(|c| &c.expansion_handle)
3962 .collect::<Vec<_>>()
3963 );
3964
3965 let conn3 = rusqlite::Connection::open_in_memory().expect("in-memory 3");
3968 crate::schema::initialize(&conn3).expect("init schema 3");
3969 insert_memory_with_embedding(
3970 &conn3,
3971 "m_b3",
3972 "cookie chocolate TOPIC_B recipe".to_string().as_str(),
3973 &emb,
3974 );
3975
3976 let bundle_noop_floor = retrieve_context_with_embedder(
3977 &conn3,
3978 "/fake-repo",
3979 &weights,
3980 ContextRequest {
3981 stage: "localization".to_string(),
3982 query: "cookie chocolate TOPIC_A".to_string(),
3984 budget_tokens: 4000,
3985 min_semantic_score: 0.0, ..Default::default()
3987 },
3988 &[],
3989 &emb,
3990 )
3991 .expect("retrieve noop floor");
3992
3993 assert!(
3995 bundle_noop_floor
3996 .capsules
3997 .iter()
3998 .any(|c| c.expansion_handle == "memory:m_b3"),
3999 "with floor=0.0 (disabled), off-topic-cosine memory must still surface via FTS; \
4000 got: {:?}",
4001 bundle_noop_floor
4002 .capsules
4003 .iter()
4004 .map(|c| &c.expansion_handle)
4005 .collect::<Vec<_>>()
4006 );
4007 }
4008
4009 #[cfg(feature = "embeddings")]
4038 #[test]
4039 fn d1f_token_economy_fewer_capsules_signal_preserved() {
4040 struct OracleTopicEmbedder;
4042 impl embeddings::Embedder for OracleTopicEmbedder {
4043 fn embed(&self, text: &str) -> Result<Vec<f32>, embeddings::EmbedderError> {
4044 let mut v = vec![0.0f32; 8];
4045 if text.contains("TOPIC_A") {
4046 v[0] = 1.0; } else {
4048 v[1] = 1.0; }
4050 Ok(v)
4051 }
4052 fn model_id(&self) -> &str {
4053 "oracle-topic-d8"
4054 }
4055 fn dim(&self) -> usize {
4056 8
4057 }
4058 }
4059
4060 let oracle = OracleTopicEmbedder;
4061
4062 let setup = |conn: &rusqlite::Connection| {
4064 for (mid, text) in [
4067 ("m_dup1", "TOPIC_A prefer ripgrep for searching"),
4068 ("m_dup2", "TOPIC_A rg is the fastest searcher"),
4069 ("m_dup3", "TOPIC_A use rg tool to find patterns"),
4070 (
4072 "m_relevant",
4073 "TOPIC_A critical lesson about search performance",
4074 ),
4075 ("m_noise1", "chocolate cookie baking TOPIC_B recipe"),
4077 ("m_noise2", "gardening tulip planting TOPIC_B spring"),
4078 ] {
4079 insert_memory_with_embedding(conn, mid, text, &oracle);
4080 }
4081 };
4082
4083 let weights = kimetsu_core::config::BrokerWeights::default();
4084
4085 let conn_lean = rusqlite::Connection::open_in_memory().expect("in-memory lean");
4092 crate::schema::initialize(&conn_lean).expect("init schema lean");
4093 setup(&conn_lean);
4094
4095 let bundle_lean = retrieve_context_with_embedder(
4096 &conn_lean,
4097 "/fake-repo",
4098 &weights,
4099 ContextRequest {
4100 stage: "localization".to_string(),
4101 query: "TOPIC_A search performance".to_string(),
4102 budget_tokens: 20_000,
4103 min_semantic_score: 0.0, ..Default::default()
4105 },
4106 &[],
4107 &embeddings::NoopEmbedder,
4108 )
4109 .expect("retrieve lean");
4110
4111 let lean_count = bundle_lean
4112 .capsules
4113 .iter()
4114 .filter(|c| c.expansion_handle.starts_with("memory:"))
4115 .count();
4116
4117 let conn_emb = rusqlite::Connection::open_in_memory().expect("in-memory emb");
4119 crate::schema::initialize(&conn_emb).expect("init schema emb");
4120 setup(&conn_emb);
4121
4122 let bundle_emb = retrieve_context_with_embedder(
4123 &conn_emb,
4124 "/fake-repo",
4125 &weights,
4126 ContextRequest {
4127 stage: "localization".to_string(),
4128 query: "TOPIC_A search performance".to_string(),
4129 budget_tokens: 20_000,
4130 min_semantic_score: 0.5, ..Default::default()
4132 },
4133 &[],
4134 &oracle,
4135 )
4136 .expect("retrieve with embeddings");
4137
4138 let emb_count = bundle_emb
4139 .capsules
4140 .iter()
4141 .filter(|c| c.expansion_handle.starts_with("memory:"))
4142 .count();
4143
4144 assert!(
4146 emb_count < lean_count,
4147 "D1e must reduce capsule count: embedding path {emb_count} must be \
4148 < lean path {lean_count}. Embedding capsules: {:?}",
4149 bundle_emb
4150 .capsules
4151 .iter()
4152 .map(|c| &c.expansion_handle)
4153 .collect::<Vec<_>>()
4154 );
4155
4156 assert!(
4158 bundle_emb
4159 .capsules
4160 .iter()
4161 .any(|c| c.expansion_handle == "memory:m_relevant"),
4162 "m_relevant must survive D1e selection (signal preserved); \
4163 embedding capsules: {:?}",
4164 bundle_emb
4165 .capsules
4166 .iter()
4167 .map(|c| &c.expansion_handle)
4168 .collect::<Vec<_>>()
4169 );
4170
4171 let lean_tokens: u32 = bundle_lean.capsules.iter().map(|c| c.token_estimate).sum();
4173 let emb_tokens: u32 = bundle_emb.capsules.iter().map(|c| c.token_estimate).sum();
4174 assert!(
4175 emb_tokens < lean_tokens,
4176 "D1e must reduce token usage: emb={emb_tokens} must be < lean={lean_tokens}"
4177 );
4178 }
4179
4180 #[test]
4186 fn lean_noop_embedder_uses_fts_then_recency_unchanged() {
4187 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4190 crate::schema::initialize(&conn).expect("init schema");
4191
4192 for (mid, text) in [
4194 ("m_x", "use git rebase to clean history"),
4195 ("m_y", "grep finds text quickly"),
4196 ] {
4197 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4198 conn.execute(
4199 "INSERT INTO memories (
4200 memory_id, scope, kind, text, normalized_text, confidence,
4201 source_event_id, provenance_snapshot_json, created_at,
4202 use_count, usefulness_score
4203 )
4204 VALUES (?1, 'global_user', 'fact', ?2, ?3, 1.0, NULL, '{}',
4205 '2026-01-01T00:00:00Z', 0, 0.0)",
4206 rusqlite::params![mid, text, normalized],
4207 )
4208 .expect("insert");
4209 conn.execute(
4210 "INSERT INTO memories_fts (memory_id, text, kind, scope) VALUES (?1, ?2, 'fact', 'global_user')",
4211 rusqlite::params![mid, text],
4212 )
4213 .expect("insert fts");
4214 }
4215
4216 let weights = kimetsu_core::config::BrokerWeights::default();
4217 let bundle = retrieve_context_with_embedder(
4219 &conn,
4220 "/fake-repo",
4221 &weights,
4222 ContextRequest {
4223 stage: "localization".to_string(),
4224 query: "grep text".to_string(),
4225 budget_tokens: 4000,
4226 ..Default::default()
4227 },
4228 &[],
4229 &embeddings::NoopEmbedder,
4230 )
4231 .expect("retrieve with NoopEmbedder must not panic");
4232
4233 let handles: Vec<&str> = bundle
4235 .capsules
4236 .iter()
4237 .filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
4238 .collect();
4239 assert!(
4240 handles.contains(&"m_y"),
4241 "m_y must surface via FTS on lean path; got {handles:?}"
4242 );
4243 }
4245
4246 #[test]
4252 fn classify_task_maps_each_kind_deterministically() {
4253 assert_eq!(
4255 classify_task("fix the panic in the parser"),
4256 TaskKind::Debug,
4257 "contains 'fix' and 'panic'"
4258 );
4259 assert_eq!(
4260 classify_task("there is a crash in auth when calling login"),
4261 TaskKind::Debug,
4262 "contains 'crash'"
4263 );
4264 assert_eq!(
4265 classify_task("debug the failing test"),
4266 TaskKind::Debug,
4267 "contains 'debug' and 'fail'"
4268 );
4269
4270 assert_eq!(
4272 classify_task("investigate why retrieval is slow"),
4273 TaskKind::Investigation,
4274 "contains 'investigate' and 'why'"
4275 );
4276 assert_eq!(
4277 classify_task("analyze the root cause of the latency"),
4278 TaskKind::Investigation,
4279 "contains 'analyze' and 'root cause'"
4280 );
4281
4282 assert_eq!(
4284 classify_task("refactor the auth module"),
4285 TaskKind::Refactor,
4286 "contains 'refactor'"
4287 );
4288 assert_eq!(
4289 classify_task("rename the config struct"),
4290 TaskKind::Refactor,
4291 "contains 'rename'"
4292 );
4293 assert_eq!(
4294 classify_task("simplify the retry handling logic"),
4295 TaskKind::Refactor,
4296 "contains 'simplify'"
4297 );
4298
4299 assert_eq!(
4301 classify_task("document the API endpoints"),
4302 TaskKind::Docs,
4303 "contains 'document'"
4304 );
4305 assert_eq!(
4306 classify_task("update the readme with new instructions"),
4307 TaskKind::Docs,
4308 "contains 'readme'"
4309 );
4310 assert_eq!(
4311 classify_task("add a docstring to the main function"),
4312 TaskKind::Docs,
4313 "contains 'docstring'"
4314 );
4315
4316 assert_eq!(
4318 classify_task("add a dark mode toggle"),
4319 TaskKind::Feature,
4320 "no debug/refactor/docs/investigate keyword"
4321 );
4322 assert_eq!(
4323 classify_task("implement the new caching layer"),
4324 TaskKind::Feature,
4325 "no debug/refactor/docs/investigate keyword"
4326 );
4327 assert_eq!(
4328 classify_task("build the export pipeline"),
4329 TaskKind::Feature,
4330 "no debug/refactor/docs/investigate keyword"
4331 );
4332 }
4333
4334 #[test]
4336 fn classify_task_respects_precedence_order() {
4337 assert_eq!(
4339 classify_task("fix and refactor the login module"),
4340 TaskKind::Debug,
4341 "Debug > Refactor"
4342 );
4343 assert_eq!(
4345 classify_task("investigate and refactor the cache layer"),
4346 TaskKind::Investigation,
4347 "Investigation > Refactor"
4348 );
4349 assert_eq!(
4351 classify_task("investigate the docs and document the API"),
4352 TaskKind::Investigation,
4353 "Investigation > Docs"
4354 );
4355 assert_eq!(
4357 classify_task("refactor and add docs"),
4358 TaskKind::Refactor,
4359 "Refactor > Docs"
4360 );
4361 assert_eq!(
4363 classify_task("fix the bug and investigate the regression"),
4364 TaskKind::Debug,
4365 "Debug > Investigation"
4366 );
4367 }
4368
4369 #[test]
4372 fn weights_for_task_kind_renormalizes_to_unit_sum() {
4373 let base = StageWeights {
4374 relevance: 0.50,
4375 confidence: 0.20,
4376 freshness: 0.20,
4377 scope: 0.10,
4378 };
4379 let original_sum = base.relevance + base.confidence + base.freshness + base.scope;
4380
4381 for kind in [
4382 TaskKind::Debug,
4383 TaskKind::Refactor,
4384 TaskKind::Investigation,
4385 TaskKind::Docs,
4386 ] {
4387 let w = weights_for_task_kind(base.clone(), kind);
4388 let new_sum = w.relevance + w.confidence + w.freshness + w.scope;
4389 assert!(
4391 (new_sum - original_sum).abs() < 1e-4,
4392 "weights_for_task_kind({kind:?}) sum {new_sum} differs from {original_sum}"
4393 );
4394 }
4395 }
4396
4397 #[test]
4399 fn weights_for_task_kind_feature_is_unchanged() {
4400 let base = StageWeights {
4401 relevance: 0.40,
4402 confidence: 0.30,
4403 freshness: 0.20,
4404 scope: 0.10,
4405 };
4406 let w = weights_for_task_kind(base.clone(), TaskKind::Feature);
4407 assert!((w.relevance - base.relevance).abs() < f32::EPSILON);
4408 assert!((w.confidence - base.confidence).abs() < f32::EPSILON);
4409 assert!((w.freshness - base.freshness).abs() < f32::EPSILON);
4410 assert!((w.scope - base.scope).abs() < f32::EPSILON);
4411 }
4412
4413 #[test]
4416 fn weights_for_task_kind_debug_up_freshness_fraction() {
4417 let base = StageWeights {
4418 relevance: 0.50,
4419 confidence: 0.20,
4420 freshness: 0.20,
4421 scope: 0.10,
4422 };
4423 let debug_w = weights_for_task_kind(base.clone(), TaskKind::Debug);
4424 assert!(
4426 debug_w.freshness > base.freshness,
4427 "Debug must increase freshness fraction: {debug_w:?}"
4428 );
4429 }
4430
4431 #[test]
4434 fn weights_for_task_kind_refactor_up_scope_fraction() {
4435 let base = StageWeights {
4436 relevance: 0.50,
4437 confidence: 0.20,
4438 freshness: 0.20,
4439 scope: 0.10,
4440 };
4441 let refactor_w = weights_for_task_kind(base.clone(), TaskKind::Refactor);
4442 assert!(
4443 refactor_w.scope > base.scope,
4444 "Refactor must increase scope fraction: {refactor_w:?}"
4445 );
4446 }
4447
4448 #[test]
4451 fn task_kind_feature_is_retrieval_neutral() {
4452 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4453 crate::schema::initialize(&conn).expect("init schema");
4454
4455 for (mid, db_kind, text) in [
4459 ("m1", "failure_pattern", "linker not found error in build"),
4460 ("m2", "convention", "use snake_case for all identifiers"),
4461 ("m3", "fact", "the cache is invalidated on every deploy"),
4462 ] {
4463 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4464 conn.execute(
4465 "INSERT INTO memories (
4466 memory_id, scope, kind, text, normalized_text, confidence,
4467 source_event_id, provenance_snapshot_json, created_at,
4468 use_count, usefulness_score
4469 )
4470 VALUES (?1, 'project', ?2, ?3, ?4, 1.0, NULL, '{}',
4471 '2026-01-01T00:00:00Z', 0, 0.0)",
4472 rusqlite::params![mid, db_kind, text, normalized],
4473 )
4474 .expect("insert memory");
4475 conn.execute(
4476 "INSERT INTO memories_fts (memory_id, text, kind, scope)
4477 VALUES (?1, ?2, ?3, 'project')",
4478 rusqlite::params![mid, text, db_kind],
4479 )
4480 .expect("insert fts");
4481 }
4482
4483 let weights = kimetsu_core::config::BrokerWeights::default();
4484 let query = "cache convention failure".to_string();
4485
4486 let baseline = retrieve_context_with_embedder(
4488 &conn,
4489 "/fake-repo",
4490 &weights,
4491 ContextRequest {
4492 stage: "localization".to_string(),
4493 query: query.clone(),
4494 budget_tokens: 4000,
4495 ..Default::default()
4496 },
4497 &[],
4498 &embeddings::NoopEmbedder,
4499 )
4500 .expect("baseline retrieve");
4501
4502 let feature = retrieve_context_with_embedder(
4504 &conn,
4505 "/fake-repo",
4506 &weights,
4507 ContextRequest {
4508 stage: "localization".to_string(),
4509 query: query.clone(),
4510 budget_tokens: 4000,
4511 task_kind: TaskKind::Feature,
4512 ..Default::default()
4513 },
4514 &[],
4515 &embeddings::NoopEmbedder,
4516 )
4517 .expect("feature retrieve");
4518
4519 let baseline_ids: Vec<&str> = baseline
4520 .capsules
4521 .iter()
4522 .map(|c| c.expansion_handle.as_str())
4523 .collect();
4524 let feature_ids: Vec<&str> = feature
4525 .capsules
4526 .iter()
4527 .map(|c| c.expansion_handle.as_str())
4528 .collect();
4529 assert_eq!(
4530 baseline_ids, feature_ids,
4531 "task_kind=Feature must produce identical retrieval to default; \
4532 baseline={baseline_ids:?} feature={feature_ids:?}"
4533 );
4534
4535 let baseline_scores: Vec<f32> = baseline.capsules.iter().map(|c| c.score).collect();
4536 let feature_scores: Vec<f32> = feature.capsules.iter().map(|c| c.score).collect();
4537 for (b, f) in baseline_scores.iter().zip(feature_scores.iter()) {
4538 assert!(
4539 (b - f).abs() < 1e-5,
4540 "scores must be identical: baseline={b} feature={f}"
4541 );
4542 }
4543 }
4544
4545 #[test]
4556 fn debug_surfaces_more_failure_pattern_than_docs() {
4557 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4558 crate::schema::initialize(&conn).expect("init schema");
4559
4560 for (i, text) in [
4564 "auth token expired causes login failure",
4565 "auth service crash on null pointer",
4566 "auth regression after upgrade breaks sessions",
4567 "auth error when certificate is invalid",
4568 ]
4569 .iter()
4570 .enumerate()
4571 {
4572 let mid = format!("mfp{i}");
4573 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4574 conn.execute(
4575 "INSERT INTO memories (
4576 memory_id, scope, kind, text, normalized_text, confidence,
4577 source_event_id, provenance_snapshot_json, created_at,
4578 use_count, usefulness_score
4579 )
4580 VALUES (?1, 'project', 'failure_pattern', ?2, ?3, 1.0, NULL, '{}',
4581 '2026-01-01T00:00:00Z', 0, 0.0)",
4582 rusqlite::params![mid, text, normalized],
4583 )
4584 .expect("insert failure_pattern");
4585 conn.execute(
4586 "INSERT INTO memories_fts (memory_id, text, kind, scope)
4587 VALUES (?1, ?2, 'failure_pattern', 'project')",
4588 rusqlite::params![mid, text],
4589 )
4590 .expect("insert fts");
4591 }
4592
4593 for (i, (db_kind, text)) in [
4596 ("convention", "auth module uses bearer tokens by convention"),
4597 ("convention", "auth scopes are documented in the API guide"),
4598 ("fact", "auth service runs on port 8443 in production"),
4599 ("fact", "auth uses JWT with RS256 signing for all tokens"),
4600 ]
4601 .iter()
4602 .enumerate()
4603 {
4604 let mid = format!("mconv{i}");
4605 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4606 conn.execute(
4607 "INSERT INTO memories (
4608 memory_id, scope, kind, text, normalized_text, confidence,
4609 source_event_id, provenance_snapshot_json, created_at,
4610 use_count, usefulness_score
4611 )
4612 VALUES (?1, 'project', ?2, ?3, ?4, 1.0, NULL, '{}',
4613 '2026-01-01T00:00:00Z', 0, 0.0)",
4614 rusqlite::params![mid, db_kind, text, normalized],
4615 )
4616 .expect("insert convention/fact");
4617 conn.execute(
4618 "INSERT INTO memories_fts (memory_id, text, kind, scope)
4619 VALUES (?1, ?2, ?3, 'project')",
4620 rusqlite::params![mid, text, db_kind],
4621 )
4622 .expect("insert fts");
4623 }
4624
4625 let weights = kimetsu_core::config::BrokerWeights::default();
4626 let query = "auth token failure".to_string();
4627
4628 let debug_bundle = retrieve_context_with_embedder(
4630 &conn,
4631 "/fake-repo",
4632 &weights,
4633 ContextRequest {
4634 stage: "localization".to_string(),
4635 query: query.clone(),
4636 budget_tokens: 4000,
4637 max_capsules: 4,
4638 task_kind: TaskKind::Debug,
4639 ..Default::default()
4640 },
4641 &[],
4642 &embeddings::NoopEmbedder,
4643 )
4644 .expect("debug retrieve");
4645
4646 let docs_bundle = retrieve_context_with_embedder(
4648 &conn,
4649 "/fake-repo",
4650 &weights,
4651 ContextRequest {
4652 stage: "localization".to_string(),
4653 query: query.clone(),
4654 budget_tokens: 4000,
4655 max_capsules: 4,
4656 task_kind: TaskKind::Docs,
4657 ..Default::default()
4658 },
4659 &[],
4660 &embeddings::NoopEmbedder,
4661 )
4662 .expect("docs retrieve");
4663
4664 let count_failure_pattern = |bundle: &ContextBundle| -> usize {
4667 bundle
4668 .capsules
4669 .iter()
4670 .filter(|c| capsule_matches_kind(c, "failure_pattern"))
4671 .count()
4672 };
4673
4674 let debug_fp = count_failure_pattern(&debug_bundle);
4675 let docs_fp = count_failure_pattern(&docs_bundle);
4676
4677 assert!(
4678 debug_fp > docs_fp,
4679 "Debug must surface strictly more failure_pattern capsules than Docs: \
4680 debug_fp={debug_fp} docs_fp={docs_fp}\n\
4681 Debug capsules: {:?}\n\
4682 Docs capsules: {:?}",
4683 debug_bundle
4684 .capsules
4685 .iter()
4686 .map(|c| format!("{}:{}", c.kind, &c.summary[..c.summary.len().min(60)]))
4687 .collect::<Vec<_>>(),
4688 docs_bundle
4689 .capsules
4690 .iter()
4691 .map(|c| format!("{}:{}", c.kind, &c.summary[..c.summary.len().min(60)]))
4692 .collect::<Vec<_>>(),
4693 );
4694 }
4695
4696 fn init_db_with_memory(memory_id: &str, text: &str) -> rusqlite::Connection {
4699 let conn = rusqlite::Connection::open_in_memory().expect("open in-memory");
4700 crate::schema::initialize(&conn).expect("init schema");
4701 let normalized = kimetsu_core::memory::normalize_memory_text(text);
4702 conn.execute(
4703 "INSERT INTO memories (
4704 memory_id, scope, kind, text, normalized_text, confidence,
4705 source_event_id, provenance_snapshot_json, created_at,
4706 use_count, usefulness_score
4707 )
4708 VALUES (?1, 'project', 'fact', ?2, ?3, 1.0, NULL, '{}',
4709 '2026-01-01T00:00:00Z', 0, 0.0)",
4710 rusqlite::params![memory_id, text, normalized],
4711 )
4712 .expect("insert memory");
4713 conn
4714 }
4715
4716 #[test]
4718 fn resolve_capsule_memory_returns_full_text() {
4719 let conn = init_db_with_memory("test-mem-id", "Use rg over grep for speed");
4720 let repo_root = std::path::Path::new("/fake-repo");
4721 let result =
4722 resolve_capsule(&conn, repo_root, "memory:test-mem-id").expect("should resolve");
4723 assert_eq!(result, "Use rg over grep for speed");
4724 }
4725
4726 #[test]
4728 fn resolve_capsule_memory_missing_id_returns_err() {
4729 let conn = init_db_with_memory("real-id", "some text");
4730 let repo_root = std::path::Path::new("/fake-repo");
4731 let err = resolve_capsule(&conn, repo_root, "memory:nonexistent-id")
4732 .expect_err("should error for missing memory");
4733 assert!(
4734 err.to_string().contains("no active memory"),
4735 "error message should mention missing: {err}"
4736 );
4737 }
4738
4739 #[test]
4741 fn resolve_capsule_file_returns_bounded_content() {
4742 let dir = make_test_dir("f2_file_resolve");
4743 let content = "hello from the file\n";
4744 std::fs::write(dir.join("notes.txt"), content).expect("write");
4745 let result = resolve_capsule(
4746 &rusqlite::Connection::open_in_memory().expect("open"),
4748 &dir,
4749 "file:notes.txt",
4750 )
4751 .expect("should resolve file");
4752 assert!(result.contains("hello from the file"));
4753 std::fs::remove_dir_all(&dir).ok();
4754 }
4755
4756 #[test]
4758 fn resolve_capsule_file_caps_large_file() {
4759 let dir = make_test_dir("f2_file_cap");
4760 let big = "A".repeat(FILE_EXPAND_CAP_BYTES * 3);
4761 std::fs::write(dir.join("big.txt"), &big).expect("write");
4762 let result = resolve_capsule(
4763 &rusqlite::Connection::open_in_memory().expect("open"),
4764 &dir,
4765 "file:big.txt",
4766 )
4767 .expect("should resolve large file");
4768 assert!(
4769 result.len() <= FILE_EXPAND_CAP_BYTES + 200,
4770 "result should be bounded: got {} bytes",
4771 result.len()
4772 );
4773 assert!(
4774 result.contains("truncated"),
4775 "truncation marker should be present"
4776 );
4777 std::fs::remove_dir_all(&dir).ok();
4778 }
4779
4780 #[test]
4782 fn resolve_capsule_unknown_handle_returns_err() {
4783 let conn = rusqlite::Connection::open_in_memory().expect("open");
4784 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "blob:abc123")
4785 .expect_err("should error");
4786 assert!(
4787 err.to_string().contains("unrecognised handle"),
4788 "got: {err}"
4789 );
4790 }
4791
4792 #[test]
4794 fn resolve_capsule_malformed_handle_returns_err() {
4795 let conn = rusqlite::Connection::open_in_memory().expect("open");
4796 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "justnocolon")
4797 .expect_err("should error");
4798 assert!(
4799 err.to_string().contains("unrecognised handle"),
4800 "got: {err}"
4801 );
4802 }
4803
4804 #[test]
4806 fn resolve_capsule_run_handle_returns_deferred_err() {
4807 let conn = rusqlite::Connection::open_in_memory().expect("open");
4808 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "run:some-run-id")
4809 .expect_err("run: should be deferred err");
4810 assert!(err.to_string().contains("not yet supported"), "got: {err}");
4811 }
4812
4813 #[test]
4815 fn resolve_capsule_file_rejects_absolute_path() {
4816 let conn = rusqlite::Connection::open_in_memory().expect("open");
4817 let err = resolve_capsule(&conn, std::path::Path::new("/r"), "file:/etc/passwd")
4818 .expect_err("should reject absolute path");
4819 assert!(err.to_string().contains("absolute path"), "got: {err}");
4820 }
4821
4822 fn make_capsule(summary: &str, score: f32) -> ContextCapsule {
4825 ContextCapsule {
4826 id: new_id().to_string(),
4827 kind: "memory".to_string(),
4828 summary: summary.to_string(),
4829 token_estimate: 10,
4830 expansion_handle: format!("memory:{}", new_id()),
4831 provenance: vec![],
4832 confidence: 1.0,
4833 freshness: 1.0,
4834 relevance: 1.0,
4835 scope_weight: 1.0,
4836 score,
4837 }
4838 }
4839
4840 #[test]
4843 fn rerank_capsules_reorders_by_query_overlap() {
4844 use crate::embeddings::StubReranker;
4845
4846 let query = "rust async tokio";
4849 let high_overlap = make_capsule("rust async tokio runtime", 0.0);
4850 let low_overlap = make_capsule("python django framework", 0.0);
4851 let capsules = vec![low_overlap.clone(), high_overlap.clone()];
4853
4854 let ranked = rerank_capsules(query, capsules, &StubReranker, 0.0, 0);
4855
4856 assert_eq!(ranked.len(), 2, "both capsules should survive (floor=0)");
4857 assert!(
4859 ranked[0].summary.contains("rust"),
4860 "rust capsule must be first, got: {:?}",
4861 ranked[0].summary
4862 );
4863 assert!(
4865 ranked[0].score > 0.05,
4866 "score must be overwritten by reranker: {}",
4867 ranked[0].score
4868 );
4869 assert!(
4871 ranked[0].score > ranked[1].score,
4872 "high overlap must score higher: {} vs {}",
4873 ranked[0].score,
4874 ranked[1].score
4875 );
4876 }
4877
4878 #[test]
4882 fn rerank_capsules_floor_drops_zero_overlap() {
4883 use crate::embeddings::StubReranker;
4884
4885 let query = "rust async tokio";
4886 let high = make_capsule("rust async tokio runtime", 0.0);
4887 let zero = make_capsule("completely unrelated document xyz", 0.0); let capsules = vec![high, zero];
4890 let ranked = rerank_capsules(query, capsules, &StubReranker, 0.3, 0);
4891
4892 assert_eq!(ranked.len(), 1, "zero-overlap capsule must be dropped");
4894 assert!(
4895 ranked[0].summary.contains("rust"),
4896 "only rust capsule should survive"
4897 );
4898 }
4899
4900 #[test]
4902 fn rerank_capsules_cap_truncates() {
4903 use crate::embeddings::StubReranker;
4904
4905 let query = "alpha beta gamma";
4906 let capsules = vec![
4907 make_capsule("alpha beta gamma delta", 0.0),
4908 make_capsule("alpha beta", 0.0),
4909 make_capsule("alpha", 0.0),
4910 make_capsule("unrelated xyz", 0.0),
4911 ];
4912
4913 let ranked = rerank_capsules(query, capsules, &StubReranker, 0.0, 2);
4914 assert_eq!(ranked.len(), 2, "cap=2 must truncate to 2 results");
4915 assert!(
4917 ranked[0].score >= ranked[1].score,
4918 "results must be sorted descending"
4919 );
4920 }
4921
4922 #[test]
4924 fn rerank_capsules_fail_open_preserves_input_order() {
4925 struct FailingReranker;
4926 impl crate::embeddings::Reranker for FailingReranker {
4927 fn rerank(
4928 &self,
4929 _query: &str,
4930 _docs: &[&str],
4931 ) -> Result<Vec<f32>, crate::embeddings::EmbedderError> {
4932 Err(crate::embeddings::EmbedderError::EmbedFailed(
4933 "simulated failure".into(),
4934 ))
4935 }
4936 fn model_id(&self) -> &str {
4937 "fail-reranker"
4938 }
4939 }
4940
4941 let query = "anything";
4942 let c1 = make_capsule("first capsule", 0.9);
4943 let c2 = make_capsule("second capsule", 0.5);
4944 let c3 = make_capsule("third capsule", 0.1);
4945 let capsules = vec![c1.clone(), c2.clone(), c3.clone()];
4946
4947 let out = rerank_capsules(query, capsules, &FailingReranker, 0.0, 0);
4948
4949 assert_eq!(out.len(), 3, "all capsules must be returned on error");
4951 assert_eq!(out[0].summary, c1.summary, "order must be preserved");
4952 assert_eq!(out[1].summary, c2.summary, "order must be preserved");
4953 assert_eq!(out[2].summary, c3.summary, "order must be preserved");
4954 }
4955
4956 #[test]
4958 fn rerank_capsules_empty_input_returns_empty() {
4959 use crate::embeddings::StubReranker;
4960 let out = rerank_capsules("query", vec![], &StubReranker, 0.0, 0);
4961 assert!(out.is_empty());
4962 }
4963
4964 #[test]
4968 fn compress_for_render_short_text_unchanged() {
4969 let text = "project:fact - Use cargo fmt before committing.";
4970 let out = compress_for_render(text, 3);
4971 assert_eq!(out, text, "short text must not be altered");
4972 }
4973
4974 #[test]
4976 fn compress_for_render_strips_tags_prefix() {
4977 let text = "[tags: rust, cargo] Always run cargo clippy before submitting a PR.";
4978 let out = compress_for_render(text, 3);
4979 assert!(
4980 !out.starts_with('['),
4981 "tags prefix must be stripped, got: {out:?}"
4982 );
4983 assert!(
4984 out.contains("cargo clippy"),
4985 "body must remain, got: {out:?}"
4986 );
4987 }
4988
4989 #[test]
4991 fn compress_for_render_strips_context_suffix() {
4992 let text =
4993 "project:fact - Use cargo fmt. Always clippy clean. (context: Kimetsu brain lesson)";
4994 let out = compress_for_render(text, 5);
4995 assert!(
4996 !out.contains("(context:"),
4997 "context suffix must be stripped, got: {out:?}"
4998 );
4999 assert!(out.contains("cargo fmt"), "body must remain, got: {out:?}");
5000 }
5001
5002 #[test]
5004 fn compress_for_render_caps_sentences() {
5005 let text =
5006 "project:fact - First sentence. Second sentence. Third sentence. Fourth sentence.";
5007 let out = compress_for_render(text, 2);
5008 assert!(out.contains("First"), "first sentence must be present");
5010 assert!(out.contains("Second"), "second sentence must be present");
5011 assert!(
5012 !out.contains("Third"),
5013 "third sentence must be truncated, got: {out:?}"
5014 );
5015 }
5016
5017 #[test]
5019 fn compress_for_render_preserves_scope_prefix() {
5020 let text = "global_user:convention - First rule. Second rule. Third rule. Fourth rule.";
5021 let out = compress_for_render(text, 2);
5022 assert!(
5023 out.starts_with("global_user:convention - "),
5024 "scope prefix must be preserved, got: {out:?}"
5025 );
5026 assert!(out.contains("First"), "first sentence must remain");
5027 assert!(!out.contains("Third"), "third sentence must be truncated");
5028 }
5029
5030 #[test]
5032 fn compress_for_render_empty_input_safe() {
5033 let out = compress_for_render("", 3);
5034 assert_eq!(out, "", "empty input must return empty string");
5035 }
5036
5037 #[test]
5039 fn compress_for_render_zero_max_sentences_returns_original() {
5040 let text = "project:fact - Some lesson that is quite long. It keeps going. And going.";
5041 let out = compress_for_render(text, 0);
5042 assert_eq!(out, text);
5043 }
5044
5045 #[test]
5047 fn compress_for_render_utf8_safe() {
5048 let text = "project:fact - こんにちは世界. Hello world. Third sentence. Fourth sentence.";
5049 let out = compress_for_render(text, 2);
5051 assert!(!out.is_empty(), "UTF-8 text must not produce empty output");
5052 assert!(!out.contains("Third"), "third sentence must be truncated");
5054 }
5055
5056 #[test]
5059 fn compress_for_render_long_memory_reduces_tokens_by_25_percent() {
5060 let long_summary = "project:fact - When a SQLite WAL file exists from a crashed process, \
5062 opening the DB causes the WAL to be replayed. The replayed WAL may contain \
5063 partial writes that corrupt the DB. Always check for WAL files before opening. \
5064 Delete the WAL only after verifying the DB is consistent. Use PRAGMA integrity_check \
5065 to validate after opening. If integrity_check fails, restore from backup. Never \
5066 truncate the WAL without replaying it first. This pattern applies to any \
5067 crash-recovery scenario.";
5068
5069 let raw_tokens = estimate_tokens(long_summary);
5070 assert!(
5071 raw_tokens > 60,
5072 "test precondition: raw memory must be >60 tokens, got {raw_tokens}"
5073 );
5074
5075 let compressed = compress_for_render(long_summary, 3);
5076 let compressed_tokens = estimate_tokens(&compressed);
5077
5078 let reduction = 1.0 - (compressed_tokens as f64 / raw_tokens as f64);
5079 assert!(
5080 reduction >= 0.25,
5081 "compression must reduce tokens by >=25% on long memories; \
5082 raw={raw_tokens} compressed={compressed_tokens} reduction={reduction:.2}"
5083 );
5084 }
5085}