tiny-agent 0.3.0

use crate::{error::StorageError, shared::Message};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};

#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct Transcript {
    pub session_id: String,
    pub messages: Vec<Message>,
    /// 父会话 id；顶层会话为 `None`。子会话由 [`crate::AgentRunTime::create_subsession`] 建立。
    pub parent: Option<String>,
    /// 本会话派生出的子会话 id 列表（父 → 子）。供"按一次 turn 下钻到各子 agent 会话"用。
    pub subsessions: Vec<String>,
}

impl Transcript {
    pub fn new(session_id: &str) -> Self {
        Self {
            session_id: String::from(session_id),
            messages: Vec::new(),
            parent: None,
            subsessions: Vec::new(),
        }
    }

    /// 带父会话 id 的子会话 transcript（`parent` 已填好）。
    pub fn with_parent(session_id: &str, parent_session_id: &str) -> Self {
        Self {
            parent: Some(parent_session_id.to_string()),
            ..Self::new(session_id)
        }
    }

    pub fn append(&mut self, message: Message) {
        self.messages.push(message);
    }
}

#[async_trait]
pub trait TranscriptStorage: Send + Sync {
    async fn save_transcript(
        &self,
        session_id: &str,
        transcript: &Transcript,
    ) -> Result<(), StorageError>;
    async fn get_transcript(&self, session_id: &str) -> Result<Option<Transcript>, StorageError>;
    async fn append_message(&self, session_id: &str, message: Message) -> Result<(), StorageError> {
        self.append_messages(session_id, vec![message]).await
    }
    async fn append_messages(
        &self,
        session_id: &str,
        messages: Vec<Message>,
    ) -> Result<(), StorageError> {
        let mut transcript = self
            .get_transcript(session_id)
            .await?
            .unwrap_or_else(|| Transcript::new(session_id));
        for message in messages {
            transcript.append(message);
        }
        self.save_transcript(session_id, &transcript).await
    }

    /// 在父会话上登记一条子会话（父 → 子）。默认实现读父 transcript、去重后追加并存回。
    ///
    /// 注意：read-modify-write，**非原子**；并发追加可能丢更新。当前 graph 节点是顺序
    /// 创建子会话的，同一 turn 内安全；若将来引入并行子 agent，需在存储层做原子 append 或加锁。
    async fn add_subsession(
        &self,
        parent_session_id: &str,
        child_session_id: &str,
    ) -> Result<(), StorageError> {
        let mut transcript = self
            .get_transcript(parent_session_id)
            .await?
            .unwrap_or_else(|| Transcript::new(parent_session_id));
        if !transcript
            .subsessions
            .iter()
            .any(|s| s == child_session_id)
        {
            transcript.subsessions.push(child_session_id.to_string());
            self.save_transcript(parent_session_id, &transcript).await?;
        }
        Ok(())
    }

    /// 读取某会话登记的子会话 id 列表（不存在则空）。
    async fn get_subsessions(&self, session_id: &str) -> Result<Vec<String>, StorageError> {
        Ok(self
            .get_transcript(session_id)
            .await?
            .map(|t| t.subsessions)
            .unwrap_or_default())
    }

    /// 删除**单个**会话的 transcript（不级联）。实现方提供这个原语。
    async fn remove_transcript(&self, session_id: &str) -> Result<(), StorageError>;

    /// 删除会话**及其全部子会话**（递归级联）：先把子树删干净，再删自身。
    /// 子会话先读出来再删自己——顺序不能反，否则读不到 subsessions。
    async fn delete_transcript(&self, session_id: &str) -> Result<(), StorageError> {
        for child in self.get_subsessions(session_id).await? {
            self.delete_transcript(&child).await?;
        }
        self.remove_transcript(session_id).await
    }
}