pe_graph/

checkpointer.rs

//! Checkpoint persistence — trait + in-memory implementation.
//!
//! The `Checkpointer` trait is storage-agnostic: it accepts and returns
//! opaque bytes. The graph engine handles serialization (bincode).
//! SurrealDB implementation lives in pe-memory (Plan 005).

use pe_core::PeError;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::SystemTime;
use tokio::sync::RwLock;

/// Metadata about a single checkpoint.
///
/// Fields may be added in future versions — construct via [`CheckpointMeta::new`].
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct CheckpointMeta {
    /// Unique checkpoint identifier.
    pub id: String,
    /// Thread this checkpoint belongs to.
    pub thread_id: String,
    /// When this checkpoint was created.
    pub created_at: SystemTime,
    /// Which superstep this was taken at.
    pub step: u32,
    /// The checkpoint this one was derived from (lineage tracking).
    /// `None` for the first checkpoint in a thread.
    pub parent_id: Option<String>,
}

impl CheckpointMeta {
    /// Create a new checkpoint metadata record.
    ///
    /// The `parent_id` defaults to `None`. Use [`with_parent`](Self::with_parent)
    /// to set lineage.
    pub fn new(id: impl Into<String>, thread_id: impl Into<String>, step: u32) -> Self {
        Self {
            id: id.into(),
            thread_id: thread_id.into(),
            created_at: SystemTime::now(),
            step,
            parent_id: None,
        }
    }

    /// Set the parent checkpoint ID (the checkpoint this one was derived from).
    pub fn with_parent(mut self, parent_id: impl Into<String>) -> Self {
        self.parent_id = Some(parent_id.into());
        self
    }
}

/// A single node's write record within a superstep.
///
/// Fields may be added in future versions — construct via [`PendingWrite::new`].
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct PendingWrite {
    /// Which node produced this write.
    pub node_name: String,
    /// Serialized update data.
    pub data: Vec<u8>,
    /// Whether the node completed successfully.
    pub success: bool,
}

impl PendingWrite {
    /// Create a new pending write record.
    pub fn new(node_name: impl Into<String>, data: Vec<u8>, success: bool) -> Self {
        Self {
            node_name: node_name.into(),
            data,
            success,
        }
    }
}

/// Storage-agnostic checkpoint persistence.
///
/// Implementations store opaque bytes — the engine handles serialization.
/// All methods are async to support network-backed stores.
#[async_trait::async_trait]
pub trait Checkpointer: Send + Sync {
    /// Save a checkpoint. Returns nothing on success.
    async fn save(
        &self,
        thread_id: &str,
        checkpoint_id: &str,
        data: &[u8],
        meta: &CheckpointMeta,
    ) -> Result<(), PeError>;

    /// Load the most recent checkpoint for a thread.
    async fn load_latest(
        &self,
        thread_id: &str,
    ) -> Result<Option<(Vec<u8>, CheckpointMeta)>, PeError>;

    /// Load a specific checkpoint by ID.
    async fn load_by_id(
        &self,
        thread_id: &str,
        checkpoint_id: &str,
    ) -> Result<Option<Vec<u8>>, PeError>;

    /// List all checkpoints for a thread, oldest first.
    async fn list(&self, thread_id: &str) -> Result<Vec<CheckpointMeta>, PeError>;

    /// Store pending writes alongside a checkpoint.
    ///
    /// **Not yet called by the engine.** The BSP loop tracks `PendingWrites`
    /// internally but does not persist them via this method yet. Plan 006
    /// (RetryPolicy) will activate this — failed nodes can be retried while
    /// successful nodes' writes are loaded from the checkpointer instead of
    /// re-executing. Implementors should store writes keyed by checkpoint_id.
    async fn put_writes(
        &self,
        thread_id: &str,
        checkpoint_id: &str,
        writes: &[PendingWrite],
    ) -> Result<(), PeError>;

    /// Delete all checkpoints for a thread.
    async fn delete_thread(&self, thread_id: &str) -> Result<(), PeError>;
}

type CheckpointEntry = (String, Vec<u8>, CheckpointMeta);

/// In-memory checkpointer for testing and short-lived graphs.
///
/// Data lives only as long as the process. For durable persistence,
/// use the SurrealDB checkpointer from pe-memory.
#[derive(Debug, Clone)]
pub struct InMemoryCheckpointer {
    store: Arc<RwLock<HashMap<String, Vec<CheckpointEntry>>>>,
}

impl InMemoryCheckpointer {
    /// Create a new empty in-memory checkpointer.
    pub fn new() -> Self {
        Self {
            store: Arc::new(RwLock::new(HashMap::new())),
        }
    }
}

impl Default for InMemoryCheckpointer {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait::async_trait]
impl Checkpointer for InMemoryCheckpointer {
    async fn save(
        &self,
        thread_id: &str,
        checkpoint_id: &str,
        data: &[u8],
        meta: &CheckpointMeta,
    ) -> Result<(), PeError> {
        let mut store = self.store.write().await;
        store.entry(thread_id.to_string()).or_default().push((
            checkpoint_id.to_string(),
            data.to_vec(),
            meta.clone(),
        ));
        Ok(())
    }

    async fn load_latest(
        &self,
        thread_id: &str,
    ) -> Result<Option<(Vec<u8>, CheckpointMeta)>, PeError> {
        let store = self.store.read().await;
        Ok(store
            .get(thread_id)
            .and_then(|entries| entries.last())
            .map(|(_, data, meta)| (data.clone(), meta.clone())))
    }

    async fn load_by_id(
        &self,
        thread_id: &str,
        checkpoint_id: &str,
    ) -> Result<Option<Vec<u8>>, PeError> {
        let store = self.store.read().await;
        Ok(store
            .get(thread_id)
            .and_then(|entries| entries.iter().find(|(id, _, _)| id == checkpoint_id))
            .map(|(_, data, _)| data.clone()))
    }

    async fn list(&self, thread_id: &str) -> Result<Vec<CheckpointMeta>, PeError> {
        let store = self.store.read().await;
        Ok(store
            .get(thread_id)
            .map(|entries| entries.iter().map(|(_, _, meta)| meta.clone()).collect())
            .unwrap_or_default())
    }

    async fn put_writes(
        &self,
        _thread_id: &str,
        _checkpoint_id: &str,
        _writes: &[PendingWrite],
    ) -> Result<(), PeError> {
        // In-memory impl: writes are already applied to state.
        // Full write tracking is for durable stores (SurrealDB, Plan 005).
        Ok(())
    }

    async fn delete_thread(&self, thread_id: &str) -> Result<(), PeError> {
        let mut store = self.store.write().await;
        store.remove(thread_id);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_meta(id: &str, thread: &str, step: u32) -> CheckpointMeta {
        CheckpointMeta::new(id, thread, step)
    }

    #[tokio::test]
    async fn test_save_and_load_latest() {
        let cp = InMemoryCheckpointer::new();
        let meta = make_meta("cp-1", "t1", 1);
        cp.save("t1", "cp-1", b"state-data", &meta).await.unwrap();

        let (data, loaded_meta) = cp.load_latest("t1").await.unwrap().unwrap();
        assert_eq!(data, b"state-data");
        assert_eq!(loaded_meta.id, "cp-1");
        assert_eq!(loaded_meta.step, 1);
    }

    #[tokio::test]
    async fn test_load_latest_returns_most_recent() {
        let cp = InMemoryCheckpointer::new();
        cp.save("t1", "cp-1", b"first", &make_meta("cp-1", "t1", 1))
            .await
            .unwrap();
        cp.save("t1", "cp-2", b"second", &make_meta("cp-2", "t1", 2))
            .await
            .unwrap();

        let (data, meta) = cp.load_latest("t1").await.unwrap().unwrap();
        assert_eq!(data, b"second");
        assert_eq!(meta.id, "cp-2");
    }

    #[tokio::test]
    async fn test_load_by_id() {
        let cp = InMemoryCheckpointer::new();
        cp.save("t1", "cp-1", b"first", &make_meta("cp-1", "t1", 1))
            .await
            .unwrap();
        cp.save("t1", "cp-2", b"second", &make_meta("cp-2", "t1", 2))
            .await
            .unwrap();

        let data = cp.load_by_id("t1", "cp-1").await.unwrap().unwrap();
        assert_eq!(data, b"first");
    }

    #[tokio::test]
    async fn test_empty_thread_returns_none() {
        let cp = InMemoryCheckpointer::new();
        assert!(cp.load_latest("nonexistent").await.unwrap().is_none());
        assert!(cp.load_by_id("nope", "nope").await.unwrap().is_none());
    }

    #[tokio::test]
    async fn test_list_checkpoints() {
        let cp = InMemoryCheckpointer::new();
        cp.save("t1", "cp-1", b"a", &make_meta("cp-1", "t1", 1))
            .await
            .unwrap();
        cp.save("t1", "cp-2", b"b", &make_meta("cp-2", "t1", 2))
            .await
            .unwrap();

        let metas = cp.list("t1").await.unwrap();
        assert_eq!(metas.len(), 2);
        assert_eq!(metas[0].id, "cp-1");
        assert_eq!(metas[1].id, "cp-2");
    }

    #[tokio::test]
    async fn test_checkpoint_meta_parent_id_default_none() {
        let meta = CheckpointMeta::new("cp-1", "t1", 1);
        assert!(meta.parent_id.is_none());
    }

    #[tokio::test]
    async fn test_checkpoint_meta_with_parent() {
        let meta = CheckpointMeta::new("cp-2", "t1", 2).with_parent("cp-1");
        assert_eq!(meta.parent_id.as_deref(), Some("cp-1"));
        assert_eq!(meta.id, "cp-2");
        assert_eq!(meta.step, 2);
    }

    #[tokio::test]
    async fn test_parent_id_preserved_through_save_load() {
        let cp = InMemoryCheckpointer::new();

        // First checkpoint — no parent
        let meta1 = make_meta("cp-1", "t1", 1);
        cp.save("t1", "cp-1", b"first", &meta1).await.unwrap();

        // Second checkpoint — parent is cp-1
        let meta2 = CheckpointMeta::new("cp-2", "t1", 2).with_parent("cp-1");
        cp.save("t1", "cp-2", b"second", &meta2).await.unwrap();

        // Load latest — should be cp-2 with parent cp-1
        let (_data, loaded_meta) = cp.load_latest("t1").await.unwrap().unwrap();
        assert_eq!(loaded_meta.id, "cp-2");
        assert_eq!(loaded_meta.parent_id.as_deref(), Some("cp-1"));

        // List — verify lineage chain
        let metas = cp.list("t1").await.unwrap();
        assert!(metas[0].parent_id.is_none()); // cp-1 has no parent
        assert_eq!(metas[1].parent_id.as_deref(), Some("cp-1")); // cp-2 -> cp-1
    }
}