diff --git a/crates/clawhdf5-agent/src/schema.rs b/crates/clawhdf5-agent/src/schema.rs index 6219bb7..a366fe3 100644 --- a/crates/clawhdf5-agent/src/schema.rs +++ b/crates/clawhdf5-agent/src/schema.rs @@ -391,20 +391,46 @@ fn load_memory_group( let tags = read_string_dataset_from_group(&group, "tags")?; let tombstones = read_u8_dataset(&group, "tombstones")?; - // Read norms if present, otherwise compute from embeddings - let norms = match read_f32_dataset(&group, "norms") { - Ok(n) if n.len() == n.len() => n, - _ => { - // Compute norms from flat embeddings - flat_embeddings - .chunks(embedding_dim) - .map(|chunk| { - let sq_sum: f32 = chunk.iter().map(|x| x * x).sum(); - sq_sum.sqrt() - }) - .collect() + // Every per-record dataset must describe exactly `n` records. Without + // this, a truncated or hand-edited file loads "successfully" and then + // panics on the first out-of-bounds index during search/delete. + if embedding_dim == 0 { + return Err(MemoryError::Schema(format!( + "/memory has {n} records but embedding_dim is 0" + ))); + } + let expected_flat = n.checked_mul(embedding_dim).ok_or_else(|| { + MemoryError::Schema(format!("/memory size overflow: {n} x {embedding_dim}")) + })?; + let check_len = |name: &str, actual: usize, expected: usize| { + if actual == expected { + Ok(()) + } else { + Err(MemoryError::Schema(format!( + "/memory/{name} has {actual} entries, expected {expected} \ + ({n} records)" + ))) } }; + check_len("embeddings", flat_embeddings.len(), expected_flat)?; + check_len("source_channel", source_channels.len(), n)?; + check_len("timestamps", timestamps.len(), n)?; + check_len("session_ids", session_ids.len(), n)?; + check_len("tags", tags.len(), n)?; + check_len("tombstones", tombstones.len(), n)?; + + // Norms are derived data: use the stored ones only if they are present + // and the right length, otherwise recompute from the embeddings. + let norms = match read_f32_dataset(&group, "norms") { + Ok(stored) if stored.len() == n => stored, + _ => flat_embeddings + .chunks(embedding_dim) + .map(|chunk| { + let sq_sum: f32 = chunk.iter().map(|x| x * x).sum(); + sq_sum.sqrt() + }) + .collect(), + }; // Unflatten embeddings let embeddings: Vec> = flat_embeddings @@ -616,3 +642,78 @@ fn read_u8_dataset(group: &clawhdf5::Group<'_>, name: &str) -> Result, M .map_err(|e| MemoryError::Hdf5(format!("cannot read u8 from {name}: {e}")))?; Ok(data.into_iter().map(|v| v as u8).collect()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn config() -> MemoryConfig { + MemoryConfig::new(std::path::PathBuf::from("unused.h5"), "agent", 4) + } + + fn cache_with(n: usize) -> MemoryCache { + let mut cache = MemoryCache::new(4); + for i in 0..n { + cache.push( + format!("chunk {i}"), + vec![i as f32 + 1.0, 0.0, 0.0, 0.0], + "user".into(), + i as f64, + "s".into(), + "t".into(), + ); + } + cache + } + + fn roundtrip(cache: &MemoryCache) -> Result { + let bytes = build_hdf5_file( + &config(), + cache, + &SessionCache::new(), + &KnowledgeCache::new(), + )?; + let file = + clawhdf5::File::from_bytes(bytes).map_err(|e| MemoryError::Hdf5(e.to_string()))?; + validate_and_load(&file).map(|(_, cache, _, _)| cache) + } + + #[test] + fn consistent_store_loads() { + let loaded = roundtrip(&cache_with(3)).unwrap(); + assert_eq!(loaded.chunks.len(), 3); + assert_eq!(loaded.norms, vec![1.0, 2.0, 3.0]); + } + + #[test] + fn wrong_length_norms_are_recomputed_not_trusted() { + // Regression: the guard used to be `n.len() == n.len()`, so a norms + // dataset of any length was accepted and corrupted every cosine score. + let mut cache = cache_with(3); + cache.norms = vec![99.0]; + let loaded = roundtrip(&cache).unwrap(); + assert_eq!(loaded.norms, vec![1.0, 2.0, 3.0]); + } + + #[test] + fn mismatched_per_record_datasets_are_schema_errors() { + type Corrupt = fn(&mut MemoryCache); + let cases: [(&str, Corrupt); 5] = [ + ("tombstones", |c| c.tombstones.truncate(1)), + ("timestamps", |c| c.timestamps.truncate(1)), + ("tags", |c| c.tags.truncate(1)), + ("session_ids", |c| c.session_ids.truncate(1)), + ("source_channel", |c| c.source_channels.truncate(1)), + ]; + for (name, corrupt) in cases { + let mut cache = cache_with(3); + corrupt(&mut cache); + match roundtrip(&cache) { + Err(MemoryError::Schema(msg)) => { + assert!(msg.contains(name), "{name}: unexpected message {msg}") + } + other => panic!("{name}: expected Schema error, got {:?}", other.map(|_| ())), + } + } + } +}