use std::path::Path;
use kimetsu_core::KimetsuResult;
use rusqlite::Connection;
use crate::embeddings::{self, Embedder, EmbedderError, encode_embedding};
use crate::project::load_project;
use crate::user_brain::open_user_brain;
#[derive(Debug, Clone)]
pub struct ScopeReport {
pub scope: &'static str,
pub opened: bool,
pub total: usize,
pub candidates: usize,
pub updated: usize,
pub failed: usize,
}
impl ScopeReport {
fn skipped(scope: &'static str) -> Self {
Self {
scope,
opened: false,
total: 0,
candidates: 0,
updated: 0,
failed: 0,
}
}
}
#[derive(Debug, Clone)]
pub struct ReindexReport {
pub project: ScopeReport,
pub user: ScopeReport,
pub embedder_model_id: String,
pub embedder_noop: bool,
}
impl ReindexReport {
pub fn updated_total(&self) -> usize {
self.project.updated + self.user.updated
}
pub fn candidates_total(&self) -> usize {
self.project.candidates + self.user.candidates
}
}
#[derive(Debug, Clone, Copy)]
pub struct ReindexOptions {
pub scope: ReindexScope,
pub dry_run: bool,
pub force: bool,
pub limit: Option<usize>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReindexScope {
Project,
User,
All,
}
impl ReindexScope {
pub fn parse(value: &str) -> Result<Self, String> {
match value.trim().to_ascii_lowercase().as_str() {
"" | "all" => Ok(Self::All),
"project" | "repo" => Ok(Self::Project),
"user" | "global" => Ok(Self::User),
other => Err(format!("unknown reindex scope `{other}`")),
}
}
}
impl Default for ReindexOptions {
fn default() -> Self {
Self {
scope: ReindexScope::All,
dry_run: false,
force: false,
limit: None,
}
}
}
pub fn reindex_all(repo_start: &Path, opts: ReindexOptions) -> KimetsuResult<ReindexReport> {
reindex_all_with_embedder(repo_start, opts, embeddings::open_default_embedder())
}
pub fn reindex_all_with_embedder(
repo_start: &Path,
opts: ReindexOptions,
embedder: &(dyn Embedder + Send + Sync),
) -> KimetsuResult<ReindexReport> {
let model_id = embedder.model_id().to_string();
let noop = embedder.is_noop();
let mut remaining = opts.limit;
let mut project_report = ScopeReport::skipped("project");
let mut user_report = ScopeReport::skipped("user");
if matches!(opts.scope, ReindexScope::Project | ReindexScope::All) {
let (_paths, _config, conn) = load_project(repo_start)?;
project_report = reindex_one_conn(&conn, "project", embedder, &opts, &mut remaining)?;
}
if matches!(opts.scope, ReindexScope::User | ReindexScope::All)
&& let Some(user_conn) = open_user_brain()?
{
user_report = reindex_one_conn(&user_conn, "user", embedder, &opts, &mut remaining)?;
}
Ok(ReindexReport {
project: project_report,
user: user_report,
embedder_model_id: model_id,
embedder_noop: noop,
})
}
const REINDEX_CHUNK: usize = 256;
fn flush_reindex_chunk(
conn: &Connection,
embedder: &(dyn Embedder + Send + Sync),
pending: &mut Vec<(String, String)>,
updated: &mut usize,
failed: &mut usize,
remaining: &mut Option<usize>,
) -> KimetsuResult<bool> {
if pending.is_empty() {
return Ok(false);
}
let texts: Vec<&str> = pending.iter().map(|(_, t)| t.as_str()).collect();
let batch_result = embedder.embed_batch(&texts);
match batch_result {
Ok(vecs) => {
for ((memory_id, _), vec) in pending.iter().zip(vecs.iter()) {
if remaining.map(|r| r == 0).unwrap_or(false) {
pending.clear();
return Ok(true); }
if vec.len() == embedder.dim() {
conn.execute(
"UPDATE memories SET embedding = ?1, embedding_model = ?2 WHERE memory_id = ?3",
rusqlite::params![encode_embedding(vec), embedder.model_id(), memory_id],
)?;
*updated += 1;
if let Some(r) = remaining {
*r = r.saturating_sub(1);
}
} else {
*failed += 1;
}
}
}
Err(_) => {
for (memory_id, text) in pending.iter() {
if remaining.map(|r| r == 0).unwrap_or(false) {
pending.clear();
return Ok(true); }
match embedder.embed(text) {
Ok(vec) if vec.len() == embedder.dim() => {
conn.execute(
"UPDATE memories SET embedding = ?1, embedding_model = ?2 WHERE memory_id = ?3",
rusqlite::params![encode_embedding(&vec), embedder.model_id(), memory_id],
)?;
*updated += 1;
if let Some(r) = remaining {
*r = r.saturating_sub(1);
}
}
Ok(_) => {
*failed += 1;
}
Err(EmbedderError::NotImplemented) => {
*failed += 1;
}
Err(_) => {
*failed += 1;
}
}
}
}
}
pending.clear();
Ok(false)
}
fn reindex_one_conn(
conn: &Connection,
scope: &'static str,
embedder: &(dyn Embedder + Send + Sync),
opts: &ReindexOptions,
remaining: &mut Option<usize>,
) -> KimetsuResult<ScopeReport> {
let total: i64 = conn.query_row(
"SELECT COUNT(*) FROM memories WHERE invalidated_at IS NULL",
[],
|row| row.get(0),
)?;
let total = total.max(0) as usize;
if embedder.is_noop() {
return Ok(ScopeReport {
scope,
opened: true,
total,
candidates: 0,
updated: 0,
failed: 0,
});
}
let model_id = embedder.model_id().to_string();
let mut stmt = if opts.force {
conn.prepare(
"
SELECT memory_id, text
FROM memories
WHERE invalidated_at IS NULL
ORDER BY created_at ASC
",
)?
} else {
conn.prepare(
"
SELECT memory_id, text
FROM memories
WHERE invalidated_at IS NULL
AND (embedding_model IS NULL OR embedding_model != ?1)
ORDER BY created_at ASC
",
)?
};
let mut rows = if opts.force {
stmt.query([])?
} else {
stmt.query(rusqlite::params![model_id])?
};
let mut candidates = 0usize;
let mut updated = 0usize;
let mut failed = 0usize;
let mut pending: Vec<(String, String)> = Vec::with_capacity(REINDEX_CHUNK);
let mut exhausted = false;
while let Some(row) = rows.next()? {
if remaining.map(|r| r == 0).unwrap_or(false) {
break;
}
candidates += 1;
let memory_id: String = row.get(0)?;
let text: String = row.get(1)?;
if opts.dry_run {
continue;
}
pending.push((memory_id, text));
let chunk_target = match *remaining {
Some(r) => REINDEX_CHUNK.min(r),
None => REINDEX_CHUNK,
};
if pending.len() >= chunk_target {
exhausted = flush_reindex_chunk(
conn,
embedder,
&mut pending,
&mut updated,
&mut failed,
remaining,
)?;
if exhausted {
break;
}
}
}
if !exhausted && !opts.dry_run {
flush_reindex_chunk(
conn,
embedder,
&mut pending,
&mut updated,
&mut failed,
remaining,
)?;
}
#[cfg(feature = "embeddings")]
if !opts.dry_run && updated > 0 {
crate::ann::invalidate_sidecar(conn);
}
Ok(ScopeReport {
scope,
opened: true,
total,
candidates,
updated,
failed,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::embeddings::{StubEmbedder, encode_embedding};
use crate::user_brain::with_user_brain_disabled;
#[test]
fn reindex_scope_parser_accepts_aliases() {
assert_eq!(
ReindexScope::parse("project").unwrap(),
ReindexScope::Project
);
assert_eq!(ReindexScope::parse("repo").unwrap(), ReindexScope::Project);
assert_eq!(ReindexScope::parse("user").unwrap(), ReindexScope::User);
assert_eq!(ReindexScope::parse("global").unwrap(), ReindexScope::User);
assert_eq!(ReindexScope::parse("all").unwrap(), ReindexScope::All);
assert_eq!(ReindexScope::parse("").unwrap(), ReindexScope::All);
assert!(ReindexScope::parse("nope").is_err());
}
#[test]
fn reindex_one_conn_backfills_null_embeddings() {
with_user_brain_disabled(|| {
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,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score
)
VALUES ('m_a', 'repo', 'fact', 'use rg', 'use rg', 1.0,
NULL, '{}', '2026-05-01T00:00:00Z', 0, 0.0)
",
[],
)
.expect("insert m_a");
conn.execute(
"
INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, embedding, embedding_model
)
VALUES ('m_b', 'repo', 'fact', 'use ripgrep', 'use ripgrep', 1.0,
NULL, '{}', '2026-05-02T00:00:00Z', 0, 0.0,
?1, 'old-model-id')
",
rusqlite::params![encode_embedding(&[0.0f32; 4])],
)
.expect("insert m_b");
let stub = StubEmbedder::new();
let mut remaining = None;
let report = reindex_one_conn(
&conn,
"project",
&stub,
&ReindexOptions::default(),
&mut remaining,
)
.expect("reindex");
assert_eq!(report.total, 2);
assert_eq!(report.candidates, 2, "both rows should be candidates");
assert_eq!(report.updated, 2, "both should be updated");
assert_eq!(report.failed, 0);
for memory_id in ["m_a", "m_b"] {
let model: String = conn
.query_row(
"SELECT embedding_model FROM memories WHERE memory_id = ?1",
rusqlite::params![memory_id],
|row| row.get(0),
)
.expect("fetch model");
assert_eq!(model, stub.model_id());
let blob: Vec<u8> = conn
.query_row(
"SELECT embedding FROM memories WHERE memory_id = ?1",
rusqlite::params![memory_id],
|row| row.get(0),
)
.expect("fetch blob");
assert_eq!(
blob.len(),
stub.dim() * 4,
"stub-d8 -> 8 floats -> 32 bytes"
);
}
});
}
#[test]
fn reindex_one_conn_dry_run_does_not_mutate() {
with_user_brain_disabled(|| {
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,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score
)
VALUES ('m_a', 'repo', 'fact', 'use rg', 'use rg', 1.0,
NULL, '{}', '2026-05-01T00:00:00Z', 0, 0.0)
",
[],
)
.expect("insert");
let stub = StubEmbedder::new();
let mut remaining = None;
let report = reindex_one_conn(
&conn,
"project",
&stub,
&ReindexOptions {
dry_run: true,
..ReindexOptions::default()
},
&mut remaining,
)
.expect("dry-run");
assert_eq!(report.candidates, 1);
assert_eq!(report.updated, 0, "dry-run must not write");
let model: Option<String> = conn
.query_row(
"SELECT embedding_model FROM memories WHERE memory_id = 'm_a'",
[],
|row| row.get(0),
)
.expect("fetch");
assert!(model.is_none(), "embedding_model should still be NULL");
});
}
#[test]
fn reindex_one_conn_with_noop_embedder_returns_zero_candidates() {
with_user_brain_disabled(|| {
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,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score
)
VALUES ('m_a', 'repo', 'fact', 'use rg', 'use rg', 1.0,
NULL, '{}', '2026-05-01T00:00:00Z', 0, 0.0)
",
[],
)
.expect("insert");
let noop = embeddings::NoopEmbedder;
let mut remaining = None;
let report = reindex_one_conn(
&conn,
"project",
&noop,
&ReindexOptions::default(),
&mut remaining,
)
.expect("noop reindex");
assert_eq!(report.total, 1);
assert_eq!(report.candidates, 0, "noop should walk zero candidates");
assert_eq!(report.updated, 0);
});
}
#[test]
fn reindex_one_conn_batches_more_than_chunk_rows() {
with_user_brain_disabled(|| {
let conn = rusqlite::Connection::open_in_memory().expect("open");
crate::schema::initialize(&conn).expect("init");
let count = 300usize;
for i in 0..count {
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score
)
VALUES (?1, 'repo', 'fact', ?2, ?3, 1.0,
NULL, '{}', '2026-05-01T00:00:00Z', 0, 0.0)",
rusqlite::params![
format!("batch-test-{i:06}"),
format!("memory text number {i}"),
format!("memory text number {i}"),
],
)
.expect("insert row");
}
let stub = StubEmbedder::new();
let mut remaining = None;
let report = reindex_one_conn(
&conn,
"project",
&stub,
&ReindexOptions::default(),
&mut remaining,
)
.expect("batch reindex");
assert_eq!(
report.candidates, count,
"all {count} rows should be candidates"
);
assert_eq!(report.updated, count, "all {count} rows should be updated");
assert_eq!(report.failed, 0, "no rows should fail with StubEmbedder");
let (check_model, check_blob_len): (String, usize) = conn
.query_row(
"SELECT embedding_model, length(embedding) FROM memories
WHERE memory_id = 'batch-test-000000'",
[],
|row| Ok((row.get(0)?, row.get(1)?)),
)
.expect("fetch spot-check row");
assert_eq!(check_model, stub.model_id());
assert_eq!(check_blob_len, stub.dim() * 4, "8 floats * 4 bytes = 32");
let null_count: i64 = conn
.query_row(
"SELECT COUNT(*) FROM memories WHERE embedding IS NULL",
[],
|row| row.get(0),
)
.expect("null count");
assert_eq!(
null_count, 0,
"no rows should have NULL embedding after reindex"
);
});
}
#[test]
fn reindex_one_conn_limit_smaller_than_chunk_is_faithful() {
with_user_brain_disabled(|| {
let conn = rusqlite::Connection::open_in_memory().expect("open");
crate::schema::initialize(&conn).expect("init");
let count = 300usize;
for i in 0..count {
conn.execute(
"INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score
)
VALUES (?1, 'repo', 'fact', ?2, ?2, 1.0,
NULL, '{}', ?3, 0, 0.0)",
rusqlite::params![
format!("limit-test-{i:06}"),
format!("memory text number {i}"),
format!("2026-05-01T00:00:{:02}Z", i % 60),
],
)
.expect("insert row");
}
let stub = StubEmbedder::new();
let mut remaining = Some(10usize);
let report = reindex_one_conn(
&conn,
"project",
&stub,
&ReindexOptions {
limit: Some(10),
..ReindexOptions::default()
},
&mut remaining,
)
.expect("limited reindex");
assert_eq!(
report.candidates, 10,
"limit must cap candidates at 10, not the full chunk"
);
assert_eq!(report.updated, 10, "exactly 10 rows updated");
assert_eq!(report.failed, 0);
assert_eq!(remaining, Some(0), "budget fully consumed");
let null_count: i64 = conn
.query_row(
"SELECT COUNT(*) FROM memories WHERE embedding IS NULL",
[],
|row| row.get(0),
)
.expect("null count");
assert_eq!(null_count, (count - 10) as i64);
});
}
#[test]
fn reindex_one_conn_force_reembeds_current_model_rows() {
with_user_brain_disabled(|| {
let conn = rusqlite::Connection::open_in_memory().expect("open");
crate::schema::initialize(&conn).expect("init");
let stub = StubEmbedder::new();
conn.execute(
"
INSERT INTO memories (
memory_id, scope, kind, text, normalized_text, confidence,
source_event_id, provenance_snapshot_json, created_at,
use_count, usefulness_score, embedding, embedding_model
)
VALUES ('m_a', 'repo', 'fact', 'use rg', 'use rg', 1.0,
NULL, '{}', '2026-05-01T00:00:00Z', 0, 0.0,
?1, ?2)
",
rusqlite::params![encode_embedding(&[0.0f32; 8]), stub.model_id()],
)
.expect("insert pre-stamped row");
let mut remaining = None;
let plain = reindex_one_conn(
&conn,
"project",
&stub,
&ReindexOptions::default(),
&mut remaining,
)
.expect("plain reindex");
assert_eq!(plain.candidates, 0);
let forced = reindex_one_conn(
&conn,
"project",
&stub,
&ReindexOptions {
force: true,
..ReindexOptions::default()
},
&mut remaining,
)
.expect("forced reindex");
assert_eq!(forced.candidates, 1);
assert_eq!(forced.updated, 1);
let blob: Vec<u8> = conn
.query_row(
"SELECT embedding FROM memories WHERE memory_id = 'm_a'",
[],
|row| row.get(0),
)
.expect("fetch blob");
assert!(
blob.iter().any(|&b| b != 0),
"force should have overwritten the zero embedding"
);
});
}
}