tiny-agent 0.3.0

//! File-backed checkpoint and transcript storage.
//!
//! Directory layout:
//! ```text
//! <root>/checkpoint/<session_id>.json
//! <root>/transcript/<session_id>.json
//! ```

use crate::{
    Agent, StorageError,
    checkpoint::CheckpointStorage,
    transcript::{Transcript, TranscriptStorage},
};
use async_trait::async_trait;
use std::path::{Path, PathBuf};
use tokio::fs;

pub struct FileStorage {
    checkpoint_dir: PathBuf,
    transcript_dir: PathBuf,
}

impl FileStorage {
    /// Creates `checkpoint/` and `transcript/` under `root`.
    pub async fn new(root: impl AsRef<Path>) -> Result<Self, StorageError> {
        let root = root.as_ref();
        let checkpoint_dir = root.join("checkpoint");
        let transcript_dir = root.join("transcript");
        fs::create_dir_all(&checkpoint_dir)
            .await
            .map_err(|e| StorageError::io(e.to_string()))?;
        fs::create_dir_all(&transcript_dir)
            .await
            .map_err(|e| StorageError::io(e.to_string()))?;
        Ok(Self {
            checkpoint_dir,
            transcript_dir,
        })
    }

    pub fn root(&self) -> &Path {
        self.checkpoint_dir.parent().unwrap_or(&self.checkpoint_dir)
    }

    fn checkpoint_path(&self, session_id: &str) -> PathBuf {
        self.checkpoint_dir.join(format!("{session_id}.json"))
    }

    fn transcript_path(&self, session_id: &str) -> PathBuf {
        self.transcript_dir.join(format!("{session_id}.json"))
    }
}

async fn write_atomic(path: &Path, bytes: &[u8]) -> Result<(), StorageError> {
    let tmp = path.with_extension("json.tmp");
    fs::write(&tmp, bytes)
        .await
        .map_err(|e| StorageError::io(e.to_string()))?;
    fs::rename(&tmp, path)
        .await
        .map_err(|e| StorageError::io(e.to_string()))?;
    Ok(())
}

async fn read_json<T: serde::de::DeserializeOwned>(path: &Path) -> Result<Option<T>, StorageError> {
    match fs::read(path).await {
        Ok(bytes) => serde_json::from_slice(&bytes)
            .map(Some)
            .map_err(|e| StorageError::serde(e.to_string())),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
        Err(e) => Err(StorageError::io(e.to_string())),
    }
}

async fn remove_if_exists(path: &Path) -> Result<(), StorageError> {
    match fs::remove_file(path).await {
        Ok(()) => Ok(()),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
        Err(e) => Err(StorageError::io(e.to_string())),
    }
}

#[async_trait]
impl CheckpointStorage for FileStorage {
    async fn save_checkpoint(&self, session_id: &str, agent: &Agent) -> Result<(), StorageError> {
        let bytes =
            serde_json::to_vec_pretty(agent).map_err(|e| StorageError::serde(e.to_string()))?;
        write_atomic(&self.checkpoint_path(session_id), &bytes).await
    }

    async fn get_checkpoint(&self, session_id: &str) -> Result<Option<Agent>, StorageError> {
        read_json(&self.checkpoint_path(session_id)).await
    }

    async fn delete_checkpoint(&self, session_id: &str) -> Result<(), StorageError> {
        remove_if_exists(&self.checkpoint_path(session_id)).await
    }
}

#[async_trait]
impl TranscriptStorage for FileStorage {
    async fn save_transcript(
        &self,
        session_id: &str,
        transcript: &Transcript,
    ) -> Result<(), StorageError> {
        let bytes = serde_json::to_vec_pretty(transcript)
            .map_err(|e| StorageError::serde(e.to_string()))?;
        write_atomic(&self.transcript_path(session_id), &bytes).await
    }

    async fn get_transcript(&self, session_id: &str) -> Result<Option<Transcript>, StorageError> {
        read_json(&self.transcript_path(session_id)).await
    }

    async fn remove_transcript(&self, session_id: &str) -> Result<(), StorageError> {
        remove_if_exists(&self.transcript_path(session_id)).await
    }
}

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

    #[tokio::test]
    async fn stores_transcripts_and_checkpoints_on_disk() {
        let root =
            std::env::temp_dir().join(format!("tiny_agent_storage_{}", uuid::Uuid::new_v4()));
        let storage = FileStorage::new(&root).await.unwrap();
        let session_id = "session-1";
        let transcript = Transcript::new(session_id);

        storage
            .save_transcript(session_id, &transcript)
            .await
            .unwrap();
        storage
            .save_checkpoint(
                session_id,
                &Agent::Interrupted(AgentState::new(Some(session_id.into()))),
            )
            .await
            .unwrap();

        assert!(storage.get_transcript(session_id).await.unwrap().is_some());
        assert!(storage.get_checkpoint(session_id).await.unwrap().is_some());

        storage.delete_checkpoint(session_id).await.unwrap();
        storage.delete_transcript(session_id).await.unwrap();

        assert!(storage.get_checkpoint(session_id).await.unwrap().is_none());
        assert!(storage.get_transcript(session_id).await.unwrap().is_none());

        let _ = std::fs::remove_dir_all(root);
    }

    #[tokio::test]
    async fn add_subsession_links_parent_to_child_and_dedupes() {
        let root =
            std::env::temp_dir().join(format!("tiny_agent_subsession_{}", uuid::Uuid::new_v4()));
        let storage = FileStorage::new(&root).await.unwrap();

        storage
            .save_transcript("parent", &Transcript::new("parent"))
            .await
            .unwrap();
        storage.add_subsession("parent", "child").await.unwrap();
        storage.add_subsession("parent", "child").await.unwrap(); // 去重:重复登记不应叠加

        assert_eq!(
            storage.get_subsessions("parent").await.unwrap(),
            vec!["child".to_string()]
        );

        // with_parent 写入的 parent 字段能读回。
        storage
            .save_transcript("child", &Transcript::with_parent("child", "parent"))
            .await
            .unwrap();
        let child = storage.get_transcript("child").await.unwrap().unwrap();
        assert_eq!(child.parent.as_deref(), Some("parent"));

        let _ = std::fs::remove_dir_all(root);
    }

    #[tokio::test]
    async fn delete_transcript_cascades_to_whole_subsession_subtree() {
        let root =
            std::env::temp_dir().join(format!("tiny_agent_cascade_{}", uuid::Uuid::new_v4()));
        let storage = FileStorage::new(&root).await.unwrap();

        // parent → child → grandchild 三层。
        for id in ["parent", "child", "grandchild"] {
            storage
                .save_transcript(id, &Transcript::new(id))
                .await
                .unwrap();
        }
        storage.add_subsession("parent", "child").await.unwrap();
        storage.add_subsession("child", "grandchild").await.unwrap();

        // 删父会话应连带删掉整棵子树。
        storage.delete_transcript("parent").await.unwrap();

        for id in ["parent", "child", "grandchild"] {
            assert!(
                storage.get_transcript(id).await.unwrap().is_none(),
                "{id} 应被级联删除"
            );
        }

        let _ = std::fs::remove_dir_all(root);
    }
}