use super::query_error;
use std::path::Path;
use velesdb_core::agent::AgentMemory;
use velesdb_core::Database;
use super::edges::{export_edges_verified, reinsert_edges, same_edge_tuples};
use super::enumeration::{reinsert_batch, scroll_page, RawFact, AGENT_COLLECTIONS};
use super::state::{CollectionProgress, MigrationLock, MigrationState, Phase};
use crate::embedder::Embedder;
pub struct RebuildSource<'a> {
pub db: &'a Database,
pub memory: &'a AgentMemory,
}
pub struct RebuildDestination<'a> {
pub db: &'a Database,
pub memory: &'a AgentMemory,
}
pub struct RebuildJournal<'a> {
pub workspace: &'a Path,
pub lock: &'a MigrationLock,
}
pub enum VectorPolicy<'a> {
Reuse,
Reembed(&'a dyn Embedder),
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct RebuildOutcome {
pub facts: u64,
pub collisions: u64,
pub edges: u64,
}
pub fn rebuild(
source: &RebuildSource<'_>,
destination: &RebuildDestination<'_>,
state: &mut MigrationState,
journal: &RebuildJournal<'_>,
policy: &VectorPolicy<'_>,
batch: usize,
) -> Result<RebuildOutcome, crate::MemoryError> {
rebuild_inner(source, destination, state, journal, policy, batch, None)
}
#[cfg(test)]
pub(crate) fn rebuild_with_stop(
source: &RebuildSource<'_>,
destination: &RebuildDestination<'_>,
state: &mut MigrationState,
journal: &RebuildJournal<'_>,
policy: &VectorPolicy<'_>,
batch: usize,
stop_after_batches: Option<u64>,
) -> Result<RebuildOutcome, crate::MemoryError> {
rebuild_inner(
source,
destination,
state,
journal,
policy,
batch,
stop_after_batches,
)
}
#[derive(Default)]
struct Run {
facts: u64,
collisions: u64,
edges: u64,
batches: u64,
stop_after_batches: Option<u64>,
}
fn rebuild_inner(
source: &RebuildSource<'_>,
destination: &RebuildDestination<'_>,
state: &mut MigrationState,
journal: &RebuildJournal<'_>,
policy: &VectorPolicy<'_>,
batch: usize,
stop_after_batches: Option<u64>,
) -> Result<RebuildOutcome, crate::MemoryError> {
if state.phase != Phase::Prepared {
return Err(query_error(format!(
"the rebuild runs strictly inside {:?}, and this journal stands at \
{:?}; a pass that ran after validation would invalidate what was \
validated",
Phase::Prepared,
state.phase
)));
}
let mut run = Run {
stop_after_batches,
..Run::default()
};
for name in AGENT_COLLECTIONS {
let step = Step {
collection: name,
policy,
batch,
};
rebuild_collection(source, destination, state, journal, &step, &mut run)?;
}
Ok(RebuildOutcome {
facts: run.facts,
collisions: run.collisions,
edges: run.edges,
})
}
struct Step<'a> {
collection: &'a str,
policy: &'a VectorPolicy<'a>,
batch: usize,
}
fn rebuild_collection(
source: &RebuildSource<'_>,
destination: &RebuildDestination<'_>,
state: &mut MigrationState,
journal: &RebuildJournal<'_>,
step: &Step<'_>,
run: &mut Run,
) -> Result<(), crate::MemoryError> {
let current = *state.progress.get(step.collection).ok_or_else(|| {
query_error(format!(
"the journal carries no progress entry for '{}'; refusing to \
invent one mid-pass",
step.collection
))
})?;
match current {
CollectionProgress::Complete => return Ok(()),
CollectionProgress::Edges => {}
CollectionProgress::Facts { cursor } => {
walk_facts(source, destination, state, journal, step, run, cursor)?;
journal_progress(state, journal, step.collection, CollectionProgress::Edges)?;
}
}
run.edges += edge_pass(source, destination, step)?;
journal_progress(
state,
journal,
step.collection,
CollectionProgress::Complete,
)
}
fn walk_facts(
source: &RebuildSource<'_>,
destination: &RebuildDestination<'_>,
state: &mut MigrationState,
journal: &RebuildJournal<'_>,
step: &Step<'_>,
run: &mut Run,
mut cursor: Option<u64>,
) -> Result<(), crate::MemoryError> {
loop {
let (facts, next) = scroll_page(source.db, step.collection, cursor, step.batch)?;
if facts.is_empty() {
return Ok(());
}
let mut pairs: Vec<(RawFact, Vec<f32>)> = Vec::with_capacity(facts.len());
for fact in facts {
let vector = vector_for(step.policy, &fact)?;
pairs.push((fact, vector));
}
let outcome = reinsert_batch(destination.db, step.collection, &pairs)?;
run.facts += outcome.inserted;
run.collisions += outcome.collisions.len() as u64;
run.batches += 1;
if run.stop_after_batches == Some(run.batches) {
return Err(query_error(format!(
"rebuild interrupted by the injected stop after {} batches; the \
destination holds this batch and the journal does not — the \
exact window a crash leaves, and what a resume replays",
run.batches
)));
}
let Some(next) = next else {
return Ok(());
};
cursor = Some(next);
journal_progress(
state,
journal,
step.collection,
CollectionProgress::Facts { cursor: Some(next) },
)?;
}
}
fn vector_for(policy: &VectorPolicy<'_>, fact: &RawFact) -> Result<Vec<f32>, crate::MemoryError> {
match policy {
VectorPolicy::Reuse => Ok(fact.source_vector.clone()),
VectorPolicy::Reembed(embedder) => {
let payload: serde_json::Value =
serde_json::from_str(&fact.payload).map_err(|err| {
query_error(format!(
"fact {} carries unreadable payload: {err}",
fact.id
))
})?;
let Some(content) = payload.get("content").and_then(serde_json::Value::as_str) else {
return Err(query_error(format!(
"fact {} carries no `content` text, so `reembed` cannot \
produce its vector; skipping it would silently drop the \
fact and re-using its old vector would mix models, so the \
pass stops here",
fact.id
)));
};
embedder
.embed(content)
.map_err(|err| query_error(format!("embedding fact {} failed: {err}", fact.id)))
}
}
}
fn edge_pass(
source: &RebuildSource<'_>,
destination: &RebuildDestination<'_>,
step: &Step<'_>,
) -> Result<u64, crate::MemoryError> {
let exported = export_edges_verified(source.memory, source.db, step.collection, step.batch)?;
let outcome = reinsert_edges(destination.memory, step.collection, &exported)?;
let back = export_edges_verified(
destination.memory,
destination.db,
step.collection,
step.batch,
)?;
same_edge_tuples(&exported, &back).map_err(|difference| {
query_error(format!(
"after reinsertion the destination's edges do not match the export \
for '{}': {difference}. Either an edge was lost, or an endpoint's \
absolute expiry passed between the export and the re-read. The \
pass is resumable: re-run it, and a mismatch that PERSISTS across \
runs is real loss",
step.collection
))
})?;
Ok(outcome.inserted)
}
fn journal_progress(
state: &mut MigrationState,
journal: &RebuildJournal<'_>,
collection: &str,
progress: CollectionProgress,
) -> Result<(), crate::MemoryError> {
state.progress.insert(collection.to_owned(), progress);
state
.write(journal.workspace, journal.lock)
.map_err(query_error)
}