enact-core 0.0.2

//! Checkpoint types for save/resume execution

use crate::kernel::{GraphId, RunId};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use svix_ksuid::KsuidLike;

/// Checkpoint - saved state of execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Checkpoint {
    /// Unique checkpoint ID
    pub id: String,
    /// Run this checkpoint belongs to
    pub run_id: RunId,
    /// Graph being executed
    pub graph_id: Option<GraphId>,
    /// Current node in execution
    pub current_node: Option<String>,
    /// State at this checkpoint
    pub state: Value,
    /// Messages history (for LLM agents)
    pub messages: Vec<MessageRecord>,
    /// Tool results collected so far
    pub tool_results: HashMap<String, Value>,
    /// Created timestamp
    pub created_at: chrono::DateTime<chrono::Utc>,
    /// Metadata
    pub metadata: HashMap<String, Value>,
}

/// Message record for checkpoint
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageRecord {
    pub role: String,
    pub content: String,
}

impl Checkpoint {
    /// Create a new checkpoint
    pub fn new(run_id: RunId) -> Self {
        Self {
            id: format!("ckpt_{}", svix_ksuid::Ksuid::new(None, None)),
            run_id,
            graph_id: None,
            current_node: None,
            state: Value::Null,
            messages: Vec::new(),
            tool_results: HashMap::new(),
            created_at: chrono::Utc::now(),
            metadata: HashMap::new(),
        }
    }

    /// Set the current state
    pub fn with_state(mut self, state: Value) -> Self {
        self.state = state;
        self
    }

    /// Set the current node
    pub fn with_node(mut self, node: impl Into<String>) -> Self {
        self.current_node = Some(node.into());
        self
    }

    /// Add a message to history
    pub fn add_message(&mut self, role: impl Into<String>, content: impl Into<String>) {
        self.messages.push(MessageRecord {
            role: role.into(),
            content: content.into(),
        });
    }

    /// Add a tool result
    pub fn add_tool_result(&mut self, tool_name: impl Into<String>, result: Value) {
        self.tool_results.insert(tool_name.into(), result);
    }

    /// Set the agent name in metadata
    pub fn with_agent_name(mut self, name: impl Into<String>) -> Self {
        self.metadata
            .insert("agent_name".to_string(), Value::String(name.into()));
        self
    }

    /// Get the agent name from metadata
    pub fn agent_name(&self) -> Option<&str> {
        self.metadata.get("agent_name").and_then(|v| v.as_str())
    }
}

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

    #[test]
    fn test_checkpoint_new() {
        let run_id = ExecutionId::new();
        let checkpoint = Checkpoint::new(run_id.clone());

        assert!(checkpoint.id.starts_with("ckpt_"));
        assert_eq!(checkpoint.run_id.as_str(), run_id.as_str());
        assert!(checkpoint.graph_id.is_none());
        assert!(checkpoint.current_node.is_none());
        assert_eq!(checkpoint.state, Value::Null);
        assert!(checkpoint.messages.is_empty());
        assert!(checkpoint.tool_results.is_empty());
        assert!(checkpoint.metadata.is_empty());
    }

    #[test]
    fn test_checkpoint_with_agent_name() {
        let run_id = ExecutionId::new();
        let checkpoint = Checkpoint::new(run_id).with_agent_name("my_agent");

        assert_eq!(checkpoint.agent_name(), Some("my_agent"));
    }

    #[test]
    fn test_checkpoint_agent_name_none_when_not_set() {
        let run_id = ExecutionId::new();
        let checkpoint = Checkpoint::new(run_id);

        assert!(checkpoint.agent_name().is_none());
    }

    #[test]
    fn test_checkpoint_builder_chain() {
        let run_id = ExecutionId::new();
        let checkpoint = Checkpoint::new(run_id)
            .with_state(Value::String("test_state".to_string()))
            .with_node("node_1")
            .with_agent_name("coder");

        assert_eq!(checkpoint.state, Value::String("test_state".to_string()));
        assert_eq!(checkpoint.current_node, Some("node_1".to_string()));
        assert_eq!(checkpoint.agent_name(), Some("coder"));
    }

    #[test]
    fn test_checkpoint_agent_name_serialization() {
        let run_id = ExecutionId::new();
        let checkpoint = Checkpoint::new(run_id).with_agent_name("assistant");

        // Serialize and deserialize
        let json = serde_json::to_string(&checkpoint).unwrap();
        let deserialized: Checkpoint = serde_json::from_str(&json).unwrap();

        assert_eq!(deserialized.agent_name(), Some("assistant"));
    }

    #[test]
    fn test_checkpoint_add_message() {
        let run_id = ExecutionId::new();
        let mut checkpoint = Checkpoint::new(run_id);

        checkpoint.add_message("user", "Hello");
        checkpoint.add_message("assistant", "Hi there!");

        assert_eq!(checkpoint.messages.len(), 2);
        assert_eq!(checkpoint.messages[0].role, "user");
        assert_eq!(checkpoint.messages[0].content, "Hello");
        assert_eq!(checkpoint.messages[1].role, "assistant");
        assert_eq!(checkpoint.messages[1].content, "Hi there!");
    }

    #[test]
    fn test_checkpoint_add_tool_result() {
        let run_id = ExecutionId::new();
        let mut checkpoint = Checkpoint::new(run_id);

        checkpoint.add_tool_result("read_file", Value::String("file contents".to_string()));

        assert_eq!(checkpoint.tool_results.len(), 1);
        assert_eq!(
            checkpoint.tool_results.get("read_file"),
            Some(&Value::String("file contents".to_string()))
        );
    }
}