fix(agent): validate per-record dataset lengths when loading a store
The norms guard was the tautology `n.len() == n.len()`, so a norms dataset of any length was trusted and corrupted every cosine score; other per-record datasets were not length-checked at all, so a truncated file loaded and then panicked on the first index. Mismatches are now MemoryError::Schema, stored norms are used only when they match the record count, and embedding_dim == 0 with records present is rejected instead of panicking in chunks(0). Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
co-authored by
Claude Fable 5.1
parent
a3f7c6fe89
commit
a9f78ca5a1
@@ -391,19 +391,45 @@ fn load_memory_group(
|
|||||||
let tags = read_string_dataset_from_group(&group, "tags")?;
|
let tags = read_string_dataset_from_group(&group, "tags")?;
|
||||||
let tombstones = read_u8_dataset(&group, "tombstones")?;
|
let tombstones = read_u8_dataset(&group, "tombstones")?;
|
||||||
|
|
||||||
// Read norms if present, otherwise compute from embeddings
|
// 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") {
|
let norms = match read_f32_dataset(&group, "norms") {
|
||||||
Ok(n) if n.len() == n.len() => n,
|
Ok(stored) if stored.len() == n => stored,
|
||||||
_ => {
|
_ => flat_embeddings
|
||||||
// Compute norms from flat embeddings
|
|
||||||
flat_embeddings
|
|
||||||
.chunks(embedding_dim)
|
.chunks(embedding_dim)
|
||||||
.map(|chunk| {
|
.map(|chunk| {
|
||||||
let sq_sum: f32 = chunk.iter().map(|x| x * x).sum();
|
let sq_sum: f32 = chunk.iter().map(|x| x * x).sum();
|
||||||
sq_sum.sqrt()
|
sq_sum.sqrt()
|
||||||
})
|
})
|
||||||
.collect()
|
.collect(),
|
||||||
}
|
|
||||||
};
|
};
|
||||||
|
|
||||||
// Unflatten embeddings
|
// Unflatten embeddings
|
||||||
@@ -616,3 +642,78 @@ fn read_u8_dataset(group: &clawhdf5::Group<'_>, name: &str) -> Result<Vec<u8>, M
|
|||||||
.map_err(|e| MemoryError::Hdf5(format!("cannot read u8 from {name}: {e}")))?;
|
.map_err(|e| MemoryError::Hdf5(format!("cannot read u8 from {name}: {e}")))?;
|
||||||
Ok(data.into_iter().map(|v| v as u8).collect())
|
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<MemoryCache, MemoryError> {
|
||||||
|
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(|_| ())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user