agent_kernel/

recovery.rs

//! State recovery and checkpoint persistence for graceful shutdown and restart.
//!
//! This module provides mechanisms to persist critical agent state before termination
//! and recover that state on startup, enabling agents to resume operation without
//! losing in-flight work or configuration.

use std::time::SystemTime;

use agent_memory::Journal;
use agent_primitives::AgentId;
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use thiserror::Error;
use tracing::{debug, warn};

use crate::lifecycle::AgentState;

/// Errors from state recovery operations.
#[derive(Debug, Error)]
pub enum RecoveryError {
    /// Journal operation failed.
    #[error("journal error: {0}")]
    JournalError(String),

    /// Serialization error.
    #[error("serialization error: {0}")]
    SerializationError(#[from] serde_json::Error),

    /// No checkpoint found for recovery.
    #[error("no checkpoint found")]
    NoCheckpoint,

    /// Checkpoint is corrupted or invalid.
    #[error("checkpoint corrupted: {0}")]
    CorruptedCheckpoint(String),

    /// I/O error.
    #[error("io error: {0}")]
    IoError(#[from] std::io::Error),
}

/// Result type for recovery operations.
pub type RecoveryResult<T> = Result<T, RecoveryError>;

/// Checkpoint containing critical agent state for recovery.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateCheckpoint {
    /// Agent identifier.
    pub agent_id: AgentId,

    /// Agent lifecycle state at checkpoint time.
    pub agent_state: AgentState,

    /// Timestamp when checkpoint was created.
    pub timestamp: SystemTime,

    /// Custom application state (opaque to recovery layer).
    pub app_state: Value,

    /// Metadata about the checkpoint.
    pub metadata: CheckpointMetadata,
}

/// Metadata about a checkpoint.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointMetadata {
    /// Version of checkpoint format.
    pub version: u32,

    /// Number of in-flight tasks at checkpoint time.
    pub in_flight_tasks: u32,

    /// Optional reason for checkpoint (e.g., "shutdown", "periodic").
    pub reason: Option<String>,

    /// Custom metadata fields.
    #[serde(default)]
    pub custom: Value,
}

impl Default for CheckpointMetadata {
    fn default() -> Self {
        Self {
            version: 1,
            in_flight_tasks: 0,
            reason: None,
            custom: Value::Null,
        }
    }
}

/// Manages checkpoint persistence and recovery.
pub struct StateRecovery {
    journal: std::sync::Arc<dyn Journal>,
    checkpoint_tag: String,
}

impl std::fmt::Debug for StateRecovery {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StateRecovery")
            .field("checkpoint_tag", &self.checkpoint_tag)
            .finish()
    }
}

impl StateRecovery {
    /// Creates a new state recovery manager.
    pub fn new(journal: std::sync::Arc<dyn Journal>) -> Self {
        Self {
            journal,
            checkpoint_tag: "checkpoint".to_string(),
        }
    }

    /// Persists a checkpoint to the journal.
    ///
    /// # Arguments
    ///
    /// * `checkpoint` - The state checkpoint to persist
    ///
    /// # Errors
    ///
    /// Returns `RecoveryError` if journal write fails or serialization fails.
    pub async fn persist_checkpoint(&self, checkpoint: &StateCheckpoint) -> RecoveryResult<()> {
        let payload = serde_json::to_vec(checkpoint)?;
        let record = agent_memory::MemoryRecord::builder(
            agent_memory::MemoryChannel::System,
            Bytes::from(payload),
        )
        .tag(&self.checkpoint_tag)
        .map_err(|e| RecoveryError::JournalError(e.to_string()))?
        .metadata("checkpoint_version", json!(checkpoint.metadata.version))
        .metadata("agent_id", json!(checkpoint.agent_id.to_string()))
        .metadata(
            "agent_state",
            json!(format!("{:?}", checkpoint.agent_state)),
        )
        .build()
        .map_err(|e| RecoveryError::JournalError(e.to_string()))?;

        self.journal
            .append(&record)
            .await
            .map_err(|e| RecoveryError::JournalError(e.to_string()))?;

        debug!(
            agent_id = %checkpoint.agent_id,
            state = ?checkpoint.agent_state,
            "checkpoint persisted"
        );

        Ok(())
    }

    /// Recovers the most recent checkpoint from the journal.
    ///
    /// # Errors
    ///
    /// Returns `RecoveryError::NoCheckpoint` if no checkpoint exists.
    /// Returns `RecoveryError::CorruptedCheckpoint` if checkpoint data is invalid.
    pub async fn recover_checkpoint(&self) -> RecoveryResult<StateCheckpoint> {
        // Retrieve recent records from journal
        let records = self
            .journal
            .tail(100)
            .await
            .map_err(|e| RecoveryError::JournalError(e.to_string()))?;

        // Find the most recent checkpoint record
        for record in records.iter().rev() {
            if record.tags().contains(&self.checkpoint_tag) {
                let checkpoint: StateCheckpoint = serde_json::from_slice(record.payload())
                    .map_err(|e| {
                        RecoveryError::CorruptedCheckpoint(format!("deserialization failed: {}", e))
                    })?;

                debug!(
                    agent_id = %checkpoint.agent_id,
                    state = ?checkpoint.agent_state,
                    "checkpoint recovered"
                );

                return Ok(checkpoint);
            }
        }

        warn!("no checkpoint found in journal");
        Err(RecoveryError::NoCheckpoint)
    }

    /// Creates a checkpoint from current agent state.
    pub fn create_checkpoint(
        agent_id: AgentId,
        agent_state: AgentState,
        app_state: Value,
        in_flight_tasks: u32,
    ) -> StateCheckpoint {
        StateCheckpoint {
            agent_id,
            agent_state,
            timestamp: SystemTime::now(),
            app_state,
            metadata: CheckpointMetadata {
                version: 1,
                in_flight_tasks,
                reason: Some("shutdown".to_string()),
                custom: Value::Null,
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use agent_memory::FileJournal;
    use std::sync::Arc;

    fn temp_path() -> std::path::PathBuf {
        let mut path = std::env::temp_dir();
        path.push(format!(
            "recovery-test-{}-{}.log",
            std::process::id(),
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        path
    }

    #[tokio::test]
    async fn checkpoint_roundtrip() {
        let path = temp_path();
        let journal = Arc::new(FileJournal::open(&path).await.unwrap());
        let recovery = StateRecovery::new(journal);

        let agent_id = AgentId::random();
        let checkpoint = StateCheckpoint {
            agent_id,
            agent_state: crate::lifecycle::AgentState::Active,
            timestamp: SystemTime::now(),
            app_state: json!({"key": "value"}),
            metadata: CheckpointMetadata {
                version: 1,
                in_flight_tasks: 5,
                reason: Some("shutdown".to_string()),
                custom: Value::Null,
            },
        };

        recovery.persist_checkpoint(&checkpoint).await.unwrap();
        let recovered = recovery.recover_checkpoint().await.unwrap();

        assert_eq!(recovered.agent_id, checkpoint.agent_id);
        assert_eq!(recovered.agent_state, checkpoint.agent_state);
        assert_eq!(recovered.app_state, checkpoint.app_state);
        assert_eq!(recovered.metadata.in_flight_tasks, 5);

        if path.exists() {
            let _ = std::fs::remove_file(path);
        }
    }

    #[tokio::test]
    async fn no_checkpoint_returns_error() {
        let path = temp_path();
        let journal = Arc::new(FileJournal::open(&path).await.unwrap());
        let recovery = StateRecovery::new(journal);

        let result = recovery.recover_checkpoint().await;
        assert!(matches!(result, Err(RecoveryError::NoCheckpoint)));

        if path.exists() {
            let _ = std::fs::remove_file(path);
        }
    }

    #[tokio::test]
    async fn multiple_checkpoints_recovers_latest() {
        let path = temp_path();
        let journal = Arc::new(FileJournal::open(&path).await.unwrap());
        let recovery = StateRecovery::new(journal);

        let agent_id = AgentId::random();

        // Persist first checkpoint
        let checkpoint1 = StateCheckpoint {
            agent_id,
            agent_state: crate::lifecycle::AgentState::Ready,
            timestamp: SystemTime::now(),
            app_state: json!({"version": 1}),
            metadata: CheckpointMetadata::default(),
        };
        recovery.persist_checkpoint(&checkpoint1).await.unwrap();

        // Small delay to ensure different timestamps
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        // Persist second checkpoint
        let checkpoint2 = StateCheckpoint {
            agent_id,
            agent_state: crate::lifecycle::AgentState::Active,
            timestamp: SystemTime::now(),
            app_state: json!({"version": 2}),
            metadata: CheckpointMetadata::default(),
        };
        recovery.persist_checkpoint(&checkpoint2).await.unwrap();

        // Should recover the latest
        let recovered = recovery.recover_checkpoint().await.unwrap();
        assert_eq!(recovered.app_state, json!({"version": 2}));

        if path.exists() {
            let _ = std::fs::remove_file(path);
        }
    }
}