use std::collections::HashMap;
use kimetsu_core::KimetsuResult;
use rusqlite::Connection;
use time::OffsetDateTime;
use time::format_description::well_known::Rfc3339;
use ulid::Ulid;
use crate::embeddings::decode_embedding;
#[derive(Debug, Clone)]
pub struct ConsolidateRow {
pub memory_id: String,
pub scope: String,
pub kind: String,
pub text: String,
pub use_count: i64,
pub usefulness_score: f32,
pub last_useful_at: Option<String>,
pub created_at: String,
pub embedding: Vec<f32>,
pub model_id: String,
}
#[derive(Debug, Clone)]
pub struct MergeCluster {
pub survivor: ConsolidateRow,
pub members: Vec<ConsolidateRow>,
}
#[derive(Debug, Default)]
pub struct ConsolidateSummary {
pub clusters_found: usize,
pub memories_merged: usize,
pub citations_reassigned: usize,
}
#[derive(Debug, Clone)]
pub struct ConsolidateOptions {
pub threshold: f32,
pub dry_run: bool,
}
impl Default for ConsolidateOptions {
fn default() -> Self {
Self {
threshold: 0.92,
dry_run: false,
}
}
}
#[derive(Debug, Clone)]
pub struct DistillOptions {
pub lo: f32,
pub hi: f32,
pub min_cluster_size: usize,
}
impl Default for DistillOptions {
fn default() -> Self {
Self {
lo: 0.75,
hi: 0.85,
min_cluster_size: 3,
}
}
}
#[derive(Debug, Clone)]
pub struct DistillCluster {
pub shared_tags: Vec<String>,
pub memories: Vec<ConsolidateRow>,
}
pub fn cosine(a: &[f32], b: &[f32]) -> f32 {
if a.len() != b.len() || a.is_empty() {
return 0.0;
}
let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
if na < f32::EPSILON || nb < f32::EPSILON {
return 0.0;
}
(dot / (na * nb)).clamp(-1.0, 1.0)
}
pub fn parse_tags(text: &str) -> Vec<String> {
let lower = text.to_ascii_lowercase();
let Some(start) = lower.find("[tags:") else {
return Vec::new();
};
let after = &text[start + 6..]; let Some(end) = after.find(']') else {
return Vec::new();
};
let tag_str = &after[..end];
let mut tags: Vec<String> = tag_str
.split(',')
.map(|t| t.trim().to_ascii_lowercase())
.filter(|t| !t.is_empty())
.collect();
tags.sort();
tags.dedup();
tags
}
struct UnionFind {
parent: Vec<usize>,
}
impl UnionFind {
fn new(n: usize) -> Self {
Self {
parent: (0..n).collect(),
}
}
fn find(&mut self, x: usize) -> usize {
if self.parent[x] != x {
self.parent[x] = self.find(self.parent[x]); }
self.parent[x]
}
fn union(&mut self, x: usize, y: usize) {
let rx = self.find(x);
let ry = self.find(y);
if rx != ry {
self.parent[ry] = rx;
}
}
}
pub fn load_embeddable_rows(
conn: &Connection,
) -> KimetsuResult<HashMap<String, Vec<ConsolidateRow>>> {
let mut stmt = conn.prepare(
"SELECT memory_id, scope, kind, text, use_count, usefulness_score,
last_useful_at, created_at, embedding, embedding_model
FROM memories
WHERE invalidated_at IS NULL
AND superseded_by IS NULL
AND valid_from IS NULL AND valid_to IS NULL
AND embedding IS NOT NULL
AND embedding_model IS NOT NULL
ORDER BY created_at DESC",
)?;
let rows = stmt.query_map([], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, String>(3)?,
row.get::<_, i64>(4)?,
row.get::<_, f64>(5)?,
row.get::<_, Option<String>>(6)?,
row.get::<_, String>(7)?,
row.get::<_, Vec<u8>>(8)?,
row.get::<_, String>(9)?,
))
})?;
let mut by_model: HashMap<String, Vec<ConsolidateRow>> = HashMap::new();
for row in rows {
let (
memory_id,
scope,
kind,
text,
use_count,
usefulness_score,
last_useful_at,
created_at,
blob,
model_id,
) = row?;
let Ok(embedding) = decode_embedding(&blob, None) else {
continue;
};
if embedding.is_empty() {
continue;
}
by_model
.entry(model_id.clone())
.or_default()
.push(ConsolidateRow {
memory_id,
scope,
kind,
text,
use_count,
usefulness_score: usefulness_score as f32,
last_useful_at,
created_at,
embedding,
model_id,
});
}
Ok(by_model)
}
fn survivor_score(row: &ConsolidateRow, recency_rank: f32) -> f32 {
let usefulness = row.usefulness_score.max(0.0);
(usefulness + 1.0) * recency_rank
}
fn parse_ts(ts: &str) -> i64 {
OffsetDateTime::parse(ts, &Rfc3339)
.map(|t| t.unix_timestamp())
.unwrap_or(0)
}
fn pick_survivor(cluster: &[usize], rows: &[ConsolidateRow]) -> usize {
let mut indexed: Vec<usize> = cluster.to_vec();
indexed.sort_by(|&a, &b| {
let ta = parse_ts(
rows[a]
.last_useful_at
.as_deref()
.unwrap_or(&rows[a].created_at),
);
let tb = parse_ts(
rows[b]
.last_useful_at
.as_deref()
.unwrap_or(&rows[b].created_at),
);
tb.cmp(&ta)
});
let n = indexed.len() as f32;
let mut best_idx = indexed[0];
let mut best_score = f32::NEG_INFINITY;
for (rank, &i) in indexed.iter().enumerate() {
let recency = 1.0 - (rank as f32) / n.max(1.0);
let score = survivor_score(&rows[i], recency);
if score > best_score {
best_score = score;
best_idx = i;
}
}
best_idx
}
pub fn find_merge_clusters(rows: &[ConsolidateRow], threshold: f32) -> Vec<MergeCluster> {
if rows.len() < 2 || !threshold.is_finite() || !(0.0..=1.0).contains(&threshold) {
return Vec::new();
}
let mut buckets = std::collections::BTreeMap::new();
for (i, row) in rows.iter().enumerate() {
if row.text.trim().is_empty() {
continue;
}
buckets
.entry((&row.scope, &row.kind, &row.model_id, row.text.trim()))
.or_insert_with(Vec::new)
.push(i);
}
let mut clusters = Vec::new();
let mut comparisons_left = 100_000usize;
for (_, mut pending) in buckets {
pending.sort_by(|&a, &b| rows[a].memory_id.cmp(&rows[b].memory_id));
while pending.len() > 1 && comparisons_left > 0 {
let survivor_idx = pick_survivor(&pending, rows);
let mut remaining = Vec::new();
let mut members = Vec::new();
for i in pending {
if i == survivor_idx {
continue;
}
if comparisons_left == 0 {
remaining.push(i);
continue;
}
comparisons_left -= 1;
if cosine(&rows[survivor_idx].embedding, &rows[i].embedding) >= threshold {
members.push(rows[i].clone());
} else {
remaining.push(i);
}
}
if !members.is_empty() {
clusters.push(MergeCluster {
survivor: rows[survivor_idx].clone(),
members,
});
}
pending = remaining;
}
}
clusters.sort_by(|a, b| a.survivor.memory_id.cmp(&b.survivor.memory_id));
clusters
}
pub fn apply_merge(
conn: &Connection,
cluster: &MergeCluster,
run_id: kimetsu_core::ids::RunId,
) -> KimetsuResult<usize> {
let events: Vec<_> = cluster
.members
.iter()
.map(|member| {
kimetsu_core::event::Event::new(
run_id,
"memory.superseded",
serde_json::json!({
"memory_id": member.memory_id,
"survivor_id": cluster.survivor.memory_id,
"use_count_delta": 0,
"score_delta": 0.0,
}),
)
})
.collect();
crate::projector::apply_events_checked(conn, &events, |c| {
let mut ids = std::collections::HashSet::new();
for planned in std::iter::once(&cluster.survivor).chain(cluster.members.iter()) {
if !ids.insert(&planned.memory_id) {
return Err("merge plan repeats a memory ID".into());
}
let current: (String, String, String, bool) = c.query_row(
"SELECT scope, kind, text, invalidated_at IS NULL AND superseded_by IS NULL
AND valid_from IS NULL AND valid_to IS NULL
FROM memories WHERE memory_id = ?1",
[&planned.memory_id],
|r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)),
)?;
if !current.3
|| current.0 != planned.scope
|| current.1 != planned.kind
|| current.2 != planned.text
|| current.0 != cluster.survivor.scope
|| current.1 != cluster.survivor.kind
|| current.2.trim() != cluster.survivor.text.trim()
{
return Err("merge plan is stale or combines distinct claims".into());
}
}
Ok(())
})?;
Ok(cluster.members.len())
}
pub fn run_consolidation(
conn: &Connection,
opts: &ConsolidateOptions,
writer: &mut impl std::io::Write,
) -> KimetsuResult<ConsolidateSummary> {
let by_model = load_embeddable_rows(conn)?;
let mut all_rows: Vec<ConsolidateRow> = by_model.into_values().flatten().collect();
all_rows.sort_by(|a, b| a.memory_id.cmp(&b.memory_id));
let clusters = find_merge_clusters(&all_rows, opts.threshold);
let mut summary = ConsolidateSummary {
clusters_found: clusters.len(),
..Default::default()
};
if clusters.is_empty() {
writeln!(
writer,
"No near-duplicate clusters found (threshold={:.2}).",
opts.threshold
)?;
return Ok(summary);
}
if opts.dry_run {
writeln!(
writer,
"Dry-run: {} cluster(s) found (threshold={:.2}):",
clusters.len(),
opts.threshold
)?;
for (i, cluster) in clusters.iter().enumerate() {
writeln!(
writer,
"\nCluster {}: SURVIVOR → {} [score={:.2} uses={}]",
i + 1,
cluster.survivor.memory_id,
cluster.survivor.usefulness_score,
cluster.survivor.use_count
)?;
writeln!(writer, " Text: {}", truncate(&cluster.survivor.text, 80))?;
for m in &cluster.members {
writeln!(
writer,
" MEMBER → {} [score={:.2} uses={}]",
m.memory_id, m.usefulness_score, m.use_count
)?;
writeln!(writer, " Text: {}", truncate(&m.text, 80))?;
}
}
return Ok(summary);
}
let run_id = kimetsu_core::ids::RunId::new();
for cluster in &clusters {
match apply_merge(conn, cluster, run_id) {
Ok(merged) => {
summary.memories_merged += merged;
}
Err(e) => {
writeln!(
writer,
"warn: merge of cluster around {} failed: {e}",
cluster.survivor.memory_id
)?;
}
}
}
writeln!(
writer,
"Consolidated {} cluster(s): {} memor{} merged.",
summary.clusters_found,
summary.memories_merged,
if summary.memories_merged == 1 {
"y"
} else {
"ies"
}
)?;
Ok(summary)
}
pub fn find_distill_clusters(
rows: &[ConsolidateRow],
opts: &DistillOptions,
) -> Vec<DistillCluster> {
let n = rows.len();
if n < opts.min_cluster_size {
return Vec::new();
}
let row_tags: Vec<Vec<String>> = rows.iter().map(|r| parse_tags(&r.text)).collect();
let mut uf = UnionFind::new(n);
for i in 0..n {
for j in (i + 1)..n {
if rows[i].model_id != rows[j].model_id {
continue;
}
let sim = cosine(&rows[i].embedding, &rows[j].embedding);
if sim < opts.lo || sim > opts.hi {
continue;
}
let shared = row_tags[i].iter().any(|t| row_tags[j].contains(t));
if shared {
uf.union(i, j);
}
}
}
let mut root_to_members: HashMap<usize, Vec<usize>> = HashMap::new();
for i in 0..n {
let root = uf.find(i);
root_to_members.entry(root).or_default().push(i);
}
let mut clusters = Vec::new();
for (_, members) in root_to_members {
if members.len() < opts.min_cluster_size {
continue;
}
let mut shared_tags: Vec<String> = row_tags[members[0]].clone();
for &i in &members[1..] {
shared_tags.retain(|t| row_tags[i].contains(t));
}
if shared_tags.is_empty() {
continue; }
let memories: Vec<ConsolidateRow> = members.iter().map(|&i| rows[i].clone()).collect();
clusters.push(DistillCluster {
shared_tags,
memories,
});
}
clusters.sort_by(|a, b| a.shared_tags.cmp(&b.shared_tags));
clusters
}
fn truncate(s: &str, max: usize) -> String {
let chars: Vec<char> = s.chars().collect();
if chars.len() <= max {
s.to_string()
} else {
format!("{}…", chars[..max].iter().collect::<String>())
}
}
#[derive(Debug, Clone, Default)]
pub struct ReflectionOptions {
pub distill_opts: DistillOptions,
pub dry_run: bool,
}
#[derive(Debug, Default)]
pub struct ReflectionSummary {
pub clusters_found: usize,
pub proposals_created: usize,
}
pub trait ModelProvider {
fn complete_text(&mut self, prompt: &str) -> Option<String>;
}
const REFLECTION_SYSTEM: &str = "You are a memory synthesizer. Given these related lessons/memories, \
synthesize ONE higher-order principle that generalizes them (2-4 sentences, \
imperative, actionable). Reply with ONLY a JSON object: \
{\"principle\": \"...\", \"tags\": [\"tag1\", \"tag2\"], \"confidence\": 0.0-1.0}";
pub fn run_reflection(
conn: &Connection,
opts: &ReflectionOptions,
model: Option<&mut dyn ModelProvider>,
writer: &mut impl std::io::Write,
) -> KimetsuResult<ReflectionSummary> {
let by_model = load_embeddable_rows(conn)?;
let mut all_rows: Vec<ConsolidateRow> = by_model.into_values().flatten().collect();
all_rows.sort_by(|a, b| a.memory_id.cmp(&b.memory_id));
let clusters = find_distill_clusters(&all_rows, &opts.distill_opts);
let mut summary = ReflectionSummary {
clusters_found: clusters.len(),
..Default::default()
};
if clusters.is_empty() {
writeln!(writer, "No reflection clusters found.")?;
return Ok(summary);
}
if opts.dry_run || model.is_none() {
writeln!(writer, "{} reflection cluster(s) found:", clusters.len())?;
for (i, cluster) in clusters.iter().enumerate() {
writeln!(
writer,
"\nCluster {} [tags: {}]:",
i + 1,
cluster.shared_tags.join(", ")
)?;
for row in &cluster.memories {
writeln!(writer, " • {}", truncate(&row.text, 80))?;
}
writeln!(
writer,
" → These {} memories could be reflected into a principle.",
cluster.memories.len()
)?;
}
return Ok(summary);
}
let model = model.unwrap(); let run_id = kimetsu_core::ids::RunId::new();
for cluster in &clusters {
let memory_texts: Vec<String> = cluster
.memories
.iter()
.map(|r| format!("- {}", r.text))
.collect();
let user_msg = memory_texts.join("\n");
let prompt = format!("{REFLECTION_SYSTEM}\n\nMemories:\n{user_msg}");
let Some(response_text) = model.complete_text(&prompt) else {
writeln!(
writer,
"warn: model call failed for cluster [{}]",
cluster.shared_tags.join(", ")
)?;
continue;
};
let Some(principle_json) = parse_reflection_json(&response_text) else {
writeln!(
writer,
"warn: could not parse reflection JSON for cluster [{}]: {response_text}",
cluster.shared_tags.join(", ")
)?;
continue;
};
let principle = principle_json
.get("principle")
.and_then(|v| v.as_str())
.unwrap_or("")
.trim()
.to_string();
if principle.is_empty() {
continue;
}
let tags = principle_json
.get("tags")
.and_then(|v| v.as_array())
.map(|a| {
a.iter()
.filter_map(|s| s.as_str())
.map(|s| s.to_string())
.collect::<Vec<_>>()
})
.unwrap_or_default();
let confidence = principle_json
.get("confidence")
.and_then(|v| v.as_f64())
.unwrap_or(0.7)
.clamp(0.0, 1.0);
let proposal_id = Ulid::new().to_string();
let source_ids: Vec<&str> = cluster
.memories
.iter()
.map(|r| r.memory_id.as_str())
.collect();
let event = kimetsu_core::event::Event::new(
run_id,
"memory.proposed",
serde_json::json!({
"proposal_id": proposal_id,
"scope": "project",
"kind": "fact",
"text": principle,
"tags": tags,
"rationale": format!(
"Reflection synthesis from {} related memories [tags: {}]",
cluster.memories.len(),
cluster.shared_tags.join(", ")
),
"proposed_confidence": confidence,
"source_event_ids": source_ids,
}),
);
match crate::projector::apply_events(conn, &[event]) {
Ok(()) => {
summary.proposals_created += 1;
writeln!(writer, "Proposed: {principle}")?;
}
Err(e) => {
writeln!(writer, "warn: failed to store reflection proposal: {e}")?;
}
}
}
Ok(summary)
}
fn parse_reflection_json(text: &str) -> Option<serde_json::Value> {
let start = text.find('{')?;
let bytes = text.as_bytes();
let mut depth = 0i32;
let mut in_string = false;
let mut escaped = false;
let mut end = None;
for (i, &b) in bytes.iter().enumerate().skip(start) {
if in_string {
if escaped {
escaped = false;
} else if b == b'\\' {
escaped = true;
} else if b == b'"' {
in_string = false;
}
} else {
match b {
b'"' => in_string = true,
b'{' => depth += 1,
b'}' => {
depth -= 1;
if depth == 0 {
end = Some(i);
break;
}
}
_ => {}
}
}
}
let json_str = &text[start..=end?];
serde_json::from_str(json_str).ok()
}
#[cfg(test)]
mod tests {
use super::*;
use rusqlite::params;
#[test]
fn cosine_same_vector_is_one() {
let v = vec![1.0f32, 0.5, -0.3];
assert!((cosine(&v, &v) - 1.0).abs() < 1e-5);
}
#[test]
fn cosine_orthogonal_is_zero() {
assert!((cosine(&[1.0f32, 0.0], &[0.0f32, 1.0]) - 0.0).abs() < 1e-5);
}
#[test]
fn cosine_opposite_is_minus_one() {
assert!((cosine(&[1.0f32, 0.0], &[-1.0f32, 0.0]) + 1.0).abs() < 1e-5);
}
#[test]
fn cosine_empty_returns_zero() {
assert_eq!(cosine(&[], &[]), 0.0);
}
#[test]
fn cosine_dim_mismatch_returns_zero() {
assert_eq!(cosine(&[1.0f32], &[1.0f32, 2.0]), 0.0);
}
#[test]
fn parse_tags_extracts_tags() {
let text = "Always use cargo fmt [tags: rust, tooling, ci]";
let tags = parse_tags(text);
assert_eq!(tags, vec!["ci", "rust", "tooling"]);
}
#[test]
fn parse_tags_no_block_returns_empty() {
assert!(parse_tags("no tags here").is_empty());
}
#[test]
fn parse_tags_case_insensitive_key() {
let text = "Something [TAGS: Rust, CI]";
let tags = parse_tags(text);
assert!(tags.contains(&"rust".to_string()));
assert!(tags.contains(&"ci".to_string()));
}
#[test]
fn parse_tags_deduplicates() {
let text = "text [tags: a, b, a]";
let tags = parse_tags(text);
assert_eq!(tags.iter().filter(|t| *t == "a").count(), 1);
}
fn make_row(id: &str, vec: Vec<f32>) -> ConsolidateRow {
ConsolidateRow {
memory_id: id.to_string(),
scope: "project".to_string(),
kind: "fact".to_string(),
text: "Identical stored claim".to_string(),
use_count: 1,
usefulness_score: 1.0,
last_useful_at: None,
created_at: "2026-01-01T00:00:00Z".to_string(),
embedding: vec,
model_id: "stub".to_string(),
}
}
#[test]
fn find_merge_clusters_identical_vectors_cluster() {
let v = vec![1.0f32, 0.0, 0.0];
let rows = vec![
make_row("a", v.clone()),
make_row("b", v.clone()),
make_row("c", v.clone()),
];
let clusters = find_merge_clusters(&rows, 0.92);
assert_eq!(clusters.len(), 1, "one cluster of identical vectors");
assert_eq!(
clusters[0].members.len(),
2,
"two members (one is survivor)"
);
}
#[test]
fn merge_preserves_scope_kind_and_distinct_claims() {
let mut base = make_row("a", vec![1.0, 0.0]);
base.text = "The development port is 4317.".into();
for variant in 0..3 {
let mut other = base.clone();
other.memory_id = "b".into();
match variant {
0 => other.scope = "user".into(),
1 => other.kind = "constraint".into(),
_ => other.text = "The production port is 4317.".into(),
}
assert!(
find_merge_clusters(&[base.clone(), other], 0.92).is_empty(),
"similarity must not erase distinct applicability or a unique claim"
);
}
}
#[test]
fn merge_similarity_chain_cannot_bridge_distant_members() {
let mut rows = vec![
make_row("a", vec![1.0, 0.0]),
make_row("b", vec![0.9396926, 0.3420201]),
make_row("c", vec![0.7660444, 0.6427876]),
];
for row in &mut rows {
row.text = "Identical claim with embedding drift".into();
}
let clusters = find_merge_clusters(&rows, 0.92);
assert!(clusters.iter().all(|c| {
c.members
.iter()
.all(|m| cosine(&c.survivor.embedding, &m.embedding) >= 0.92)
}));
}
#[test]
fn find_merge_clusters_orthogonal_no_clusters() {
let rows = vec![
make_row("a", vec![1.0f32, 0.0]),
make_row("b", vec![0.0f32, 1.0]),
];
let clusters = find_merge_clusters(&rows, 0.92);
assert!(clusters.is_empty(), "orthogonal vectors do not cluster");
}
#[test]
fn find_merge_clusters_different_models_do_not_cluster() {
let v = vec![1.0f32, 0.0];
let mut r1 = make_row("a", v.clone());
r1.model_id = "model-a".to_string();
let mut r2 = make_row("b", v.clone());
r2.model_id = "model-b".to_string();
let clusters = find_merge_clusters(&[r1, r2], 0.92);
assert!(clusters.is_empty(), "different models must not cluster");
}
#[test]
fn survivor_is_highest_usefulness_score() {
let v = vec![1.0f32, 0.0, 0.0];
let mut high = make_row("high", v.clone());
high.usefulness_score = 10.0;
high.use_count = 5;
let mut low = make_row("low", v.clone());
low.usefulness_score = 0.1;
low.use_count = 1;
let clusters = find_merge_clusters(&[low, high], 0.92);
assert_eq!(clusters.len(), 1);
assert_eq!(clusters[0].survivor.memory_id, "high");
assert_eq!(clusters[0].members[0].memory_id, "low");
}
#[test]
fn find_distill_clusters_requires_shared_tags() {
let v1 = vec![1.0f32, 0.5, 0.0];
let v2 = vec![1.0f32, 0.4, 0.1];
let mut r1 = make_row("a", v1);
r1.text = "first memory [tags: rust]".to_string();
let mut r2 = make_row("b", v2);
r2.text = "second memory [tags: python]".to_string();
let mut r3 = make_row("c", vec![1.0f32, 0.4, 0.05]);
r3.text = "third memory [tags: go]".to_string();
let opts = DistillOptions {
lo: 0.7,
hi: 0.99,
min_cluster_size: 2,
};
let clusters = find_distill_clusters(&[r1, r2, r3], &opts);
assert!(clusters.is_empty(), "no shared tags → no distill cluster");
}
#[test]
fn find_distill_clusters_shared_tag_and_band_clusters() {
let v = vec![1.0f32, 0.5, 0.1];
let make = |id: &str, extra: f32| {
let mut r = make_row(id, vec![1.0 + extra, 0.5, 0.1]);
r.text = format!("memory {id} [tags: rust, ci]");
r
};
let rows = vec![make("a", 0.0), make("b", 0.001), make("c", 0.002)];
let _ = v; let opts = DistillOptions {
lo: 0.0,
hi: 1.0,
min_cluster_size: 3,
};
let clusters = find_distill_clusters(&rows, &opts);
assert!(!clusters.is_empty(), "shared tag + band → distill cluster");
assert!(
clusters[0].shared_tags.contains(&"ci".to_string()),
"shared_tags contains 'ci'"
);
}
#[test]
fn apply_merge_preserves_evidence_without_counting_copies_as_independent() {
use kimetsu_core::ids::RunId;
let conn = rusqlite::Connection::open_in_memory().expect("open");
crate::schema::initialize(&conn).expect("init");
for (id, use_count, score) in [("survivor", 3i64, 5.0f64), ("member", 2i64, 2.0f64)] {
conn.execute(
"INSERT INTO memories
(memory_id, scope, kind, text, normalized_text, confidence,
provenance_snapshot_json, created_at, use_count, usefulness_score)
VALUES (?1,'project','fact',?2,?2,0.9,'{}','2026-01-01T00:00:00Z',?3,?4)",
params![id, "Identical stored claim", use_count, score],
)
.expect("insert");
}
let survivor = ConsolidateRow {
memory_id: "survivor".to_string(),
scope: "project".to_string(),
kind: "fact".to_string(),
text: "Identical stored claim".to_string(),
use_count: 3,
usefulness_score: 5.0,
last_useful_at: None,
created_at: "2026-01-01T00:00:00Z".to_string(),
embedding: vec![1.0, 0.0],
model_id: "stub".to_string(),
};
let member = ConsolidateRow {
memory_id: "member".to_string(),
scope: "project".to_string(),
kind: "fact".to_string(),
text: "Identical stored claim".to_string(),
use_count: 2,
usefulness_score: 2.0,
last_useful_at: None,
created_at: "2026-01-01T00:00:00Z".to_string(),
embedding: vec![1.0, 0.0],
model_id: "stub".to_string(),
};
let cluster = MergeCluster {
survivor,
members: vec![member],
};
let run_id = RunId::new();
let merged = apply_merge(&conn, &cluster, run_id).expect("apply_merge");
assert_eq!(merged, 1);
let (use_count, score): (i64, f64) = conn
.query_row(
"SELECT use_count, usefulness_score FROM memories WHERE memory_id = 'survivor'",
[],
|r| Ok((r.get(0)?, r.get(1)?)),
)
.expect("query survivor");
assert_eq!(
use_count, 3,
"copying a claim cannot create independent observations"
);
assert!(
(score - 5.0).abs() < 0.01,
"keep survivor's evidence, got {score}"
);
let superseded_by: Option<String> = conn
.query_row(
"SELECT superseded_by FROM memories WHERE memory_id = 'member'",
[],
|r| r.get(0),
)
.expect("query member");
assert_eq!(superseded_by.as_deref(), Some("survivor"));
}
#[test]
fn stale_merge_plan_does_not_retire_any_member() {
let conn = rusqlite::Connection::open_in_memory().unwrap();
crate::schema::initialize(&conn).unwrap();
for id in ["a", "b", "c"] {
conn.execute("INSERT INTO memories (memory_id,scope,kind,text,normalized_text,confidence,
provenance_snapshot_json,created_at) VALUES (?1,'project','fact','Identical stored claim',
'identical stored claim',0.5,'{}','2026-01-01T00:00:00Z')", params![id]).unwrap();
}
let cluster = MergeCluster {
survivor: make_row("a", vec![1.0, 0.0]),
members: vec![make_row("b", vec![1.0, 0.0]), make_row("c", vec![1.0, 0.0])],
};
conn.execute(
"UPDATE memories SET text='A corrected distinct claim' WHERE memory_id='c'",
[],
)
.unwrap();
assert!(apply_merge(&conn, &cluster, kimetsu_core::ids::RunId::new()).is_err());
let retired: i64 = conn
.query_row(
"SELECT COUNT(*) FROM memories WHERE superseded_by IS NOT NULL",
[],
|r| r.get(0),
)
.unwrap();
assert_eq!(retired, 0, "the entire stale plan must roll back");
conn.execute("UPDATE memories SET text='Identical stored claim', valid_from='2099-01-01T00:00:00Z' WHERE memory_id='c'", []).unwrap();
assert!(
apply_merge(&conn, &cluster, kimetsu_core::ids::RunId::new()).is_err(),
"identical text with different applicability must not consolidate"
);
}
#[test]
fn citations_reassigned_on_merge() {
use kimetsu_core::ids::RunId;
let conn = rusqlite::Connection::open_in_memory().expect("open");
crate::schema::initialize(&conn).expect("init");
for id in ["survivor", "member"] {
conn.execute(
"INSERT INTO memories
(memory_id, scope, kind, text, normalized_text, confidence,
provenance_snapshot_json, created_at, use_count, usefulness_score)
VALUES (?1,'project','fact',?2,?2,0.9,'{}','2026-01-01T00:00:00Z',1,1.0)",
params![id, "Identical stored claim"],
)
.expect("insert memory");
}
conn.execute(
"INSERT INTO memory_citations (run_id, memory_id, turn, cited_at)
VALUES ('run-1', 'member', 1, '2026-01-01T00:00:00Z')",
[],
)
.expect("insert citation");
let cluster = MergeCluster {
survivor: ConsolidateRow {
memory_id: "survivor".to_string(),
scope: "project".to_string(),
kind: "fact".to_string(),
text: "Identical stored claim".to_string(),
use_count: 1,
usefulness_score: 1.0,
last_useful_at: None,
created_at: "2026-01-01T00:00:00Z".to_string(),
embedding: vec![1.0, 0.0],
model_id: "stub".to_string(),
},
members: vec![ConsolidateRow {
memory_id: "member".to_string(),
scope: "project".to_string(),
kind: "fact".to_string(),
text: "Identical stored claim".to_string(),
use_count: 1,
usefulness_score: 1.0,
last_useful_at: None,
created_at: "2026-01-01T00:00:00Z".to_string(),
embedding: vec![1.0, 0.0],
model_id: "stub".to_string(),
}],
};
apply_merge(&conn, &cluster, RunId::new()).expect("apply_merge");
let mid: String = conn
.query_row(
"SELECT memory_id FROM memory_citations WHERE run_id = 'run-1' AND turn = 1",
[],
|r| r.get(0),
)
.expect("query citation");
assert_eq!(mid, "survivor", "citation reassigned to survivor");
let member_count: i64 = conn
.query_row(
"SELECT COUNT(*) FROM memory_citations WHERE memory_id = 'member'",
[],
|r| r.get(0),
)
.expect("count member citations");
assert_eq!(member_count, 0, "member citations deleted");
}
#[test]
fn superseded_row_excluded_from_latest_memory_candidates() {
use crate::context::retrieve_context_with_embedder;
use crate::embeddings::NoopEmbedder;
use kimetsu_core::config::BrokerWeights;
let conn = rusqlite::Connection::open_in_memory().expect("open");
crate::schema::initialize(&conn).expect("init");
conn.execute(
"INSERT INTO memories
(memory_id, scope, kind, text, normalized_text, confidence,
provenance_snapshot_json, created_at, use_count, usefulness_score)
VALUES ('surv','project','fact','rust tooling','rust tooling',0.9,'{}',
'2026-01-01T00:00:00Z',1,1.0)",
[],
)
.expect("insert survivor");
conn.execute(
"INSERT INTO memories
(memory_id, scope, kind, text, normalized_text, confidence,
provenance_snapshot_json, created_at, use_count, usefulness_score,
superseded_by)
VALUES ('dup','project','fact','rust tooling dup','rust tooling dup',0.9,'{}',
'2026-01-01T00:00:00Z',1,1.0,'surv')",
[],
)
.expect("insert superseded");
conn.execute(
"INSERT INTO memories_fts (memory_id, text, kind, scope)
VALUES ('surv', 'rust tooling', 'fact', 'project')",
[],
)
.expect("insert fts");
let weights = BrokerWeights::default();
let req = crate::context::ContextRequest {
stage: "test".to_string(),
query: "rust tooling".to_string(),
budget_tokens: 4096,
..Default::default()
};
let embedder = NoopEmbedder;
let bundle = retrieve_context_with_embedder(&conn, "", &weights, req, &[], &embedder)
.expect("retrieve");
let ids: Vec<&str> = bundle
.capsules
.iter()
.chain(bundle.excluded.iter())
.filter_map(|c| c.expansion_handle.strip_prefix("memory:"))
.collect();
assert!(
!ids.contains(&"dup"),
"superseded memory must not appear in retrieval"
);
}
#[test]
fn v2_brain_migrates_to_v3_with_backup_and_superseded_by_column() {
use crate::migrate;
use kimetsu_core::KIMETSU_SCHEMA_VERSION;
let tmp_id = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
let tmp_dir = std::env::temp_dir().join(format!("kimetsu-test-v3mig-{tmp_id}"));
std::fs::create_dir_all(&tmp_dir).expect("create tmp dir");
let db_path = tmp_dir.join("brain.db");
{
let conn = rusqlite::Connection::open(&db_path).expect("open");
crate::schema::create_baseline_for_test(&conn).expect("baseline");
crate::schema::migrate_v1_to_v2(&conn).expect("v1→v2");
conn.execute(
"UPDATE schema_info SET value = 2 WHERE key = 'kimetsu_schema_version'",
[],
)
.expect("stamp v2");
conn.execute(
"INSERT INTO memories
(memory_id, scope, kind, text, normalized_text, confidence,
provenance_snapshot_json, created_at, use_count, usefulness_score)
VALUES ('m1','project','fact','hello','hello',0.9,'{}','2026-01-01T00:00:00Z',0,0.0)",
[],
).expect("insert memory");
}
{
let conn = rusqlite::Connection::open(&db_path).expect("reopen");
let outcome = migrate::run_migrations(&conn).expect("run_migrations");
assert_eq!(outcome.from, 2);
assert_eq!(outcome.to, KIMETSU_SCHEMA_VERSION);
assert!(
outcome.applied.contains(&3),
"v3 must be in applied list, got: {:?}",
outcome.applied
);
assert!(
outcome.backup_path.is_some(),
"backup must be created for non-empty brain during migration"
);
let has_superseded_by: bool = conn.query_row(
"SELECT COUNT(*) FROM pragma_table_info('memories') WHERE name = 'superseded_by'",
[],
|r| r.get::<_, i64>(0),
).map(|n| n > 0).unwrap_or(false);
assert!(
has_superseded_by,
"superseded_by column must exist after v3 migration"
);
let has_edges: bool = conn
.query_row(
"SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='memory_edges'",
[],
|r| r.get::<_, i64>(0),
)
.map(|n| n > 0)
.unwrap_or(false);
assert!(
has_edges,
"memory_edges table must exist after v4 migration"
);
}
let _ = std::fs::remove_dir_all(&tmp_dir);
}
#[test]
fn consolidation_is_rebuild_safe() {
use crate::projector;
use kimetsu_core::ids::RunId;
let conn = rusqlite::Connection::open_in_memory().expect("open");
crate::schema::initialize(&conn).expect("init");
let run_id = RunId::new();
projector::apply_events(
&conn,
&[kimetsu_core::event::Event::new(
run_id,
"run.started",
serde_json::json!({"project_id": "test", "task": "rebuild-safety"}),
)],
)
.expect("run.started");
for (mid, text) in [
("survivor", "Identical stored claim"),
("member", "Identical stored claim"),
] {
projector::apply_events(
&conn,
&[kimetsu_core::event::Event::new(
run_id,
"memory.accepted",
serde_json::json!({
"memory_id": mid,
"scope": "project",
"kind": "fact",
"text": text,
"normalized_text": text,
"confidence": 0.9
}),
)],
)
.expect("accepted");
}
conn.execute(
"UPDATE memories SET use_count = 3, usefulness_score = 5.0 \
WHERE memory_id = 'survivor'",
[],
)
.expect("seed survivor stats");
conn.execute(
"UPDATE memories SET use_count = 2, usefulness_score = 2.0 \
WHERE memory_id = 'member'",
[],
)
.expect("seed member stats");
projector::apply_events(
&conn,
&[kimetsu_core::event::Event::new(
run_id,
"memory.cited",
serde_json::json!({
"memory_id": "member",
"turn": 1,
"rationale": "test citation"
}),
)],
)
.expect("memory.cited");
let cluster = MergeCluster {
survivor: ConsolidateRow {
memory_id: "survivor".to_string(),
scope: "project".to_string(),
kind: "fact".to_string(),
text: "Identical stored claim".to_string(),
use_count: 3,
usefulness_score: 5.0,
last_useful_at: None,
created_at: "2026-01-01T00:00:00Z".to_string(),
embedding: vec![1.0, 0.0],
model_id: "stub".to_string(),
},
members: vec![ConsolidateRow {
memory_id: "member".to_string(),
scope: "project".to_string(),
kind: "fact".to_string(),
text: "Identical stored claim".to_string(),
use_count: 2,
usefulness_score: 2.0,
last_useful_at: None,
created_at: "2026-01-01T00:00:00Z".to_string(),
embedding: vec![1.0, 0.0],
model_id: "stub".to_string(),
}],
};
apply_merge(&conn, &cluster, RunId::new()).expect("apply_merge");
let (pre_uc, pre_score): (i64, f64) = conn
.query_row(
"SELECT use_count, usefulness_score FROM memories \
WHERE memory_id = 'survivor'",
[],
|r| Ok((r.get(0)?, r.get(1)?)),
)
.expect("query survivor after consolidation");
let pre_cited: String = conn
.query_row(
"SELECT memory_id FROM memory_citations WHERE turn = 1",
[],
|r| r.get(0),
)
.expect("citation must exist post-consolidation");
assert_eq!(
pre_cited, "survivor",
"pre-rebuild: citation must point at survivor"
);
projector::rebuild_in_place(&conn).expect("rebuild_in_place");
let (post_uc, post_score): (i64, f64) = conn
.query_row(
"SELECT use_count, usefulness_score FROM memories \
WHERE memory_id = 'survivor'",
[],
|r| Ok((r.get(0)?, r.get(1)?)),
)
.expect("query survivor after rebuild");
assert_eq!(post_uc, 0, "copies cannot manufacture replay evidence");
assert_eq!(post_score, 0.0);
let post_cited: String = conn
.query_row(
"SELECT memory_id FROM memory_citations WHERE turn = 1",
[],
|r| r.get(0),
)
.expect("citation must still exist after rebuild");
assert_eq!(
post_cited, "survivor",
"post-rebuild: citation must still point at survivor (got {post_cited:?})"
);
assert_eq!(pre_uc, 3);
assert_eq!(pre_score, 5.0);
}
struct MockReflector {
response: Option<String>,
}
impl ModelProvider for MockReflector {
fn complete_text(&mut self, _prompt: &str) -> Option<String> {
self.response.take()
}
}
fn insert_reflectable(conn: &rusqlite::Connection, id: &str, text: &str) {
use crate::embeddings::{Embedder, StubEmbedder, encode_embedding};
let stub = StubEmbedder::new();
let vec = stub.embed(text).expect("embed");
let blob = encode_embedding(&vec);
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, embedding, embedding_model
) VALUES (?1, 'project', 'fact', ?2, ?2, 1.0, NULL, '{}',
'2026-01-01T00:00:00Z', 0, 0.0, ?3, ?4)",
params![id, text, blob, stub.model_id()],
)
.expect("insert reflectable");
}
#[test]
fn run_reflection_creates_proposal_from_cluster() {
let conn = rusqlite::Connection::open_in_memory().expect("open");
crate::schema::initialize(&conn).expect("init");
insert_reflectable(
&conn,
"a",
"always run cargo fmt before commit [tags: rust, ci]",
);
insert_reflectable(
&conn,
"b",
"always run cargo fmt before push [tags: rust, ci]",
);
insert_reflectable(&conn, "c", "always run cargo fmt on save [tags: rust, ci]");
let mut model = MockReflector {
response: Some(
r#"{"principle": "Always format Rust code with cargo fmt before sharing.", "tags": ["rust", "ci"], "confidence": 0.85}"#
.to_string(),
),
};
let opts = ReflectionOptions {
distill_opts: DistillOptions {
lo: 0.0,
hi: 1.0,
min_cluster_size: 3,
},
dry_run: false,
};
let mut out: Vec<u8> = Vec::new();
let summary =
run_reflection(&conn, &opts, Some(&mut model), &mut out).expect("run_reflection");
assert!(
summary.clusters_found >= 1,
"must find at least one cluster"
);
assert_eq!(
summary.proposals_created, 1,
"model-backed reflection must create exactly one proposal"
);
let (text, status): (String, String) = conn
.query_row(
"SELECT text, status FROM memory_proposals LIMIT 1",
[],
|r| Ok((r.get(0)?, r.get(1)?)),
)
.expect("proposal row exists");
assert!(text.contains("cargo fmt"), "proposal carries the principle");
assert_eq!(
status, "pending",
"reflection proposal must be pending review"
);
}
#[test]
fn run_reflection_without_model_reports_only() {
let conn = rusqlite::Connection::open_in_memory().expect("open");
crate::schema::initialize(&conn).expect("init");
insert_reflectable(
&conn,
"a",
"prefer thiserror in libraries [tags: rust, errors]",
);
insert_reflectable(
&conn,
"b",
"prefer thiserror for library crates [tags: rust, errors]",
);
insert_reflectable(
&conn,
"c",
"use thiserror not anyhow in libs [tags: rust, errors]",
);
let opts = ReflectionOptions {
distill_opts: DistillOptions {
lo: 0.0,
hi: 1.0,
min_cluster_size: 3,
},
dry_run: false,
};
let mut out: Vec<u8> = Vec::new();
let summary = run_reflection(&conn, &opts, None, &mut out).expect("run_reflection");
assert!(summary.clusters_found >= 1);
assert_eq!(
summary.proposals_created, 0,
"no model → no proposals (graceful degradation)"
);
let report = String::from_utf8(out).unwrap();
assert!(
report.contains("could be reflected"),
"report must describe reflectable clusters, got: {report}"
);
let proposal_count: i64 = conn
.query_row("SELECT COUNT(*) FROM memory_proposals", [], |r| r.get(0))
.unwrap();
assert_eq!(proposal_count, 0, "no proposals written without a model");
}
}