Skip to main content

velesdb_core/agent/
snapshot.rs

1//! Snapshot and versioning support for `AgentMemory`.
2//!
3//! Provides serialization/deserialization of `AgentMemory` state for:
4//! - Persistence across restarts
5//! - Rollback to previous versions
6//! - State transfer between instances
7//!
8//! # Snapshot Format
9//!
10//! ```text
11//! [Magic: "VAMM" 4 bytes]
12//! [Version: 1 byte]
13//! [Semantic state length: 8 bytes]
14//! [Semantic state: N bytes]
15//! [Episodic state length: 8 bytes]
16//! [Episodic state: N bytes]
17//! [Procedural state length: 8 bytes]
18//! [Procedural state: N bytes]
19//! [TTL state length: 8 bytes]
20//! [TTL state: N bytes]
21//! [CRC32: 4 bytes]
22//! ```
23
24// Reason: Numeric casts in snapshot handling are intentional:
25// - usize to u32 in CRC32: i ranges 0-255, always fits in u32
26// - u64 to usize for lengths: Snapshot data is created/loaded on same architecture
27//   or architecture-compatible data. Lengths are validated before use.
28// All length values are bounds-checked against data.len() before array access.
29#![allow(clippy::cast_possible_truncation)]
30
31use std::fs::File;
32use std::io::{self, Read, Write};
33use std::path::Path;
34
35use crate::storage::snapshot::crc32_hash;
36
37/// Snapshot file magic bytes for `AgentMemory`.
38pub const SNAPSHOT_MAGIC: &[u8; 4] = b"VAMM";
39
40/// Current snapshot format version.
41///
42/// Bumped to 2 when the TTL entry layout gained a 1-byte `MemoryKind` tag
43/// (24 -> 25 bytes/entry). A v1 snapshot written by an earlier release is now
44/// rejected with [`SnapshotError::UnsupportedVersion`] instead of silently
45/// loading with every TTL dropped.
46pub const SNAPSHOT_VERSION: u8 = 2;
47
48/// Memory state for serialization.
49#[derive(Debug, Clone, Default)]
50pub struct MemoryState {
51    /// Serialized semantic memory entries.
52    pub semantic: Vec<u8>,
53    /// Serialized episodic memory entries.
54    pub episodic: Vec<u8>,
55    /// Serialized procedural memory entries.
56    pub procedural: Vec<u8>,
57    /// Serialized TTL state.
58    pub ttl: Vec<u8>,
59}
60
61/// Snapshot metadata.
62#[derive(Debug, Clone)]
63pub struct SnapshotMetadata {
64    /// Snapshot format version.
65    pub version: u8,
66    /// Total size in bytes.
67    pub total_size: usize,
68    /// CRC32 checksum.
69    pub checksum: u32,
70}
71
72/// Error type for snapshot operations.
73#[derive(Debug)]
74#[non_exhaustive]
75pub enum SnapshotError {
76    /// IO error during read/write.
77    Io(io::Error),
78    /// Invalid magic bytes.
79    InvalidMagic,
80    /// Unsupported version.
81    UnsupportedVersion(u8),
82    /// CRC checksum mismatch.
83    ChecksumMismatch {
84        /// Expected CRC32 value stored in the snapshot.
85        expected: u32,
86        /// Actual CRC32 value computed from the data.
87        actual: u32,
88    },
89    /// Data corruption or truncation.
90    CorruptedData(String),
91}
92
93impl std::fmt::Display for SnapshotError {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        match self {
96            Self::Io(e) => write!(f, "IO error: {e}"),
97            Self::InvalidMagic => write!(f, "Invalid snapshot magic bytes"),
98            Self::UnsupportedVersion(v) => write!(f, "Unsupported snapshot version: {v}"),
99            Self::ChecksumMismatch { expected, actual } => {
100                write!(
101                    f,
102                    "Checksum mismatch: expected {expected:08x}, got {actual:08x}"
103                )
104            }
105            Self::CorruptedData(msg) => write!(f, "Corrupted data: {msg}"),
106        }
107    }
108}
109
110impl std::error::Error for SnapshotError {
111    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
112        match self {
113            Self::Io(e) => Some(e),
114            _ => None,
115        }
116    }
117}
118
119impl From<io::Error> for SnapshotError {
120    fn from(e: io::Error) -> Self {
121        Self::Io(e)
122    }
123}
124
125/// Creates a snapshot from memory state.
126///
127/// # Arguments
128///
129/// * `state` - Memory state to serialize
130///
131/// # Returns
132///
133/// Serialized snapshot bytes.
134#[must_use]
135pub fn create_snapshot(state: &MemoryState) -> Vec<u8> {
136    let total_size = 4
137        + 1
138        + 8
139        + state.semantic.len()
140        + 8
141        + state.episodic.len()
142        + 8
143        + state.procedural.len()
144        + 8
145        + state.ttl.len()
146        + 4;
147    let mut buf = Vec::with_capacity(total_size);
148
149    buf.extend_from_slice(SNAPSHOT_MAGIC);
150    buf.push(SNAPSHOT_VERSION);
151
152    buf.extend_from_slice(&(state.semantic.len() as u64).to_le_bytes());
153    buf.extend_from_slice(&state.semantic);
154
155    buf.extend_from_slice(&(state.episodic.len() as u64).to_le_bytes());
156    buf.extend_from_slice(&state.episodic);
157
158    buf.extend_from_slice(&(state.procedural.len() as u64).to_le_bytes());
159    buf.extend_from_slice(&state.procedural);
160
161    buf.extend_from_slice(&(state.ttl.len() as u64).to_le_bytes());
162    buf.extend_from_slice(&state.ttl);
163
164    let crc = crc32_hash(&buf);
165    buf.extend_from_slice(&crc.to_le_bytes());
166
167    buf
168}
169
170/// Loads a snapshot from bytes.
171///
172/// # Arguments
173///
174/// * `data` - Snapshot bytes
175///
176/// # Errors
177///
178/// Returns error if snapshot is invalid or corrupted.
179pub fn load_snapshot(data: &[u8]) -> Result<MemoryState, SnapshotError> {
180    validate_snapshot_header(data)?;
181
182    let mut offset = 5; // skip magic (4) + version (1)
183    let payload_end = data.len() - 4; // exclude trailing CRC
184
185    let semantic = read_section(data, &mut offset, payload_end, "Semantic")?;
186    let episodic = read_section(data, &mut offset, payload_end, "Episodic")?;
187    let procedural = read_section(data, &mut offset, payload_end, "Procedural")?;
188    let ttl = read_section(data, &mut offset, payload_end, "TTL")?;
189
190    Ok(MemoryState {
191        semantic,
192        episodic,
193        procedural,
194        ttl,
195    })
196}
197
198/// Validates magic bytes, version, and CRC32 checksum of a snapshot.
199fn validate_snapshot_header(data: &[u8]) -> Result<(), SnapshotError> {
200    const MIN_SIZE: usize = 4 + 1 + 8 + 8 + 8 + 8 + 4;
201
202    if data.len() < MIN_SIZE {
203        return Err(SnapshotError::CorruptedData(
204            "Snapshot too small".to_string(),
205        ));
206    }
207    if &data[0..4] != SNAPSHOT_MAGIC {
208        return Err(SnapshotError::InvalidMagic);
209    }
210    let version = data[4];
211    if version != SNAPSHOT_VERSION {
212        return Err(SnapshotError::UnsupportedVersion(version));
213    }
214
215    let stored_crc = u32::from_le_bytes(
216        data[data.len() - 4..]
217            .try_into()
218            .map_err(|_| SnapshotError::CorruptedData("Invalid CRC bytes".to_string()))?,
219    );
220    let computed_crc = crc32_hash(&data[..data.len() - 4]);
221    if stored_crc != computed_crc {
222        return Err(SnapshotError::ChecksumMismatch {
223            expected: stored_crc,
224            actual: computed_crc,
225        });
226    }
227    Ok(())
228}
229
230/// Reads a length-prefixed section from the snapshot data.
231fn read_section(
232    data: &[u8],
233    offset: &mut usize,
234    payload_end: usize,
235    label: &str,
236) -> Result<Vec<u8>, SnapshotError> {
237    let section_len = read_u64(&data[*offset..])? as usize;
238    *offset += 8;
239    // Checked: a forged/corrupt `section_len` near `usize::MAX` must not wrap
240    // `*offset + section_len` (which could spuriously pass the bound and then
241    // panic on the slice). Mirrors `validate_binary_header`'s checked arithmetic.
242    let end = offset
243        .checked_add(section_len)
244        .filter(|end| *end <= payload_end)
245        .ok_or_else(|| SnapshotError::CorruptedData(format!("{label} data truncated")))?;
246    let section = data[*offset..end].to_vec();
247    *offset = end;
248    Ok(section)
249}
250
251/// Saves a snapshot to a file.
252///
253/// Uses atomic write (temp file + rename) for safety.
254///
255/// # Errors
256///
257/// Returns error if file operations fail.
258pub fn save_snapshot_to_file<P: AsRef<Path>>(
259    path: P,
260    state: &MemoryState,
261) -> Result<(), SnapshotError> {
262    let path = path.as_ref();
263    let snapshot_data = create_snapshot(state);
264
265    let temp_path = path.with_extension("tmp");
266    let mut file = File::create(&temp_path)?;
267    file.write_all(&snapshot_data)?;
268    file.sync_all()?;
269    drop(file);
270
271    std::fs::rename(&temp_path, path)?;
272
273    Ok(())
274}
275
276/// Loads a snapshot from a file.
277///
278/// # Errors
279///
280/// Returns error if file operations fail or snapshot is invalid.
281pub fn load_snapshot_from_file<P: AsRef<Path>>(path: P) -> Result<MemoryState, SnapshotError> {
282    let mut file = File::open(path)?;
283    let mut data = Vec::new();
284    file.read_to_end(&mut data)?;
285    load_snapshot(&data)
286}
287
288/// Helper to read u64 from bytes.
289fn read_u64(data: &[u8]) -> Result<u64, SnapshotError> {
290    if data.len() < 8 {
291        return Err(SnapshotError::CorruptedData(
292            "Not enough bytes for u64".to_string(),
293        ));
294    }
295    Ok(u64::from_le_bytes(data[0..8].try_into().map_err(|_| {
296        SnapshotError::CorruptedData("Invalid u64 bytes".to_string())
297    })?))
298}
299
300/// Snapshot manager for versioned snapshots.
301pub struct SnapshotManager {
302    /// Base directory for snapshots.
303    base_path: std::path::PathBuf,
304    /// Maximum number of snapshots to retain.
305    max_snapshots: usize,
306}
307
308impl SnapshotManager {
309    /// Creates a new snapshot manager.
310    ///
311    /// # Arguments
312    ///
313    /// * `base_path` - Directory for storing snapshots
314    /// * `max_snapshots` - Maximum number of snapshots to retain
315    pub fn new<P: AsRef<Path>>(base_path: P, max_snapshots: usize) -> Self {
316        Self {
317            base_path: base_path.as_ref().to_path_buf(),
318            max_snapshots,
319        }
320    }
321
322    /// Creates a new versioned snapshot.
323    ///
324    /// # Returns
325    ///
326    /// The version number of the created snapshot.
327    ///
328    /// # Errors
329    ///
330    /// Returns error if file operations fail.
331    pub fn create_versioned_snapshot(&self, state: &MemoryState) -> Result<u64, SnapshotError> {
332        std::fs::create_dir_all(&self.base_path)?;
333
334        let version = self.next_version()?;
335        let filename = format!("snapshot_{version:08}.vamm");
336        let path = self.base_path.join(filename);
337
338        save_snapshot_to_file(&path, state)?;
339        self.cleanup_old_snapshots()?;
340
341        Ok(version)
342    }
343
344    /// Loads the latest snapshot.
345    ///
346    /// # Errors
347    ///
348    /// Returns error if no snapshots exist or loading fails.
349    pub fn load_latest(&self) -> Result<(u64, MemoryState), SnapshotError> {
350        let version = self
351            .latest_version()?
352            .ok_or_else(|| SnapshotError::CorruptedData("No snapshots found".to_string()))?;
353        let state = self.load_version(version)?;
354        Ok((version, state))
355    }
356
357    /// Loads a specific snapshot version.
358    ///
359    /// # Errors
360    ///
361    /// Returns error if version doesn't exist or loading fails.
362    pub fn load_version(&self, version: u64) -> Result<MemoryState, SnapshotError> {
363        let filename = format!("snapshot_{version:08}.vamm");
364        let path = self.base_path.join(filename);
365        load_snapshot_from_file(&path)
366    }
367
368    /// Lists all available snapshot versions.
369    ///
370    /// # Errors
371    ///
372    /// Returns error if directory operations fail.
373    pub fn list_versions(&self) -> Result<Vec<u64>, SnapshotError> {
374        if !self.base_path.exists() {
375            return Ok(Vec::new());
376        }
377
378        let mut versions: Vec<u64> = std::fs::read_dir(&self.base_path)?
379            .filter_map(Result::ok)
380            .filter_map(|e| parse_snapshot_version(&e.file_name().to_string_lossy()))
381            .collect();
382
383        versions.sort_unstable();
384        Ok(versions)
385    }
386
387    /// Returns the latest snapshot version.
388    fn latest_version(&self) -> Result<Option<u64>, SnapshotError> {
389        Ok(self.list_versions()?.into_iter().max())
390    }
391
392    /// Returns the next version number.
393    fn next_version(&self) -> Result<u64, SnapshotError> {
394        Ok(self.latest_version()?.map_or(1, |v| v + 1))
395    }
396
397    /// Removes old snapshots beyond the retention limit.
398    fn cleanup_old_snapshots(&self) -> Result<(), SnapshotError> {
399        let versions = self.list_versions()?;
400        if versions.len() <= self.max_snapshots {
401            return Ok(());
402        }
403
404        let to_remove = versions.len() - self.max_snapshots;
405        for version in versions.into_iter().take(to_remove) {
406            let filename = format!("snapshot_{version:08}.vamm");
407            let path = self.base_path.join(filename);
408            let _ = std::fs::remove_file(path);
409        }
410
411        Ok(())
412    }
413}
414
415/// Extracts a snapshot version number from a filename like `snapshot_00000001.vamm`.
416fn parse_snapshot_version(filename: &str) -> Option<u64> {
417    filename
418        .strip_prefix("snapshot_")
419        .and_then(|s| s.strip_suffix(".vamm"))
420        .and_then(|s| s.parse::<u64>().ok())
421}