tiny-agent 0.2.0

use super::{Agent, AgentFailure, AgentRunTime, FailureKind, storage_failure};
use crate::{
    error::AgentError,
    observability,
    shared::{Message, Role},
    state::AgentState,
    transcript::Transcript,
};
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use uuid::Uuid;

impl AgentRunTime {
    /// 驱动一个 `Agent` 跑到停车点(终态或中断),并负责 **checkpoint 生命周期**:
    ///
    /// - `Success` / `Fail` 是终态,运行结束时清掉该会话残留的 checkpoint —— 既回收存储,
    ///   也避免下次 [`resume`](Self::resume) 读到陈旧快照。
    /// - `Interrupted` 的 checkpoint 由内层在停车时写入,这里**保留**,留给 `resume`。
    pub async fn run(
        &self,
        agent_state: Agent,
        cancellation_token: CancellationToken,
    ) -> Result<Agent, AgentError> {
        let session_id = self.get_session_id(&agent_state);
        let span = tracing::info_span!(
            "agent.run",
            "gen_ai.conversation.id" = %session_id,
            "tiny_agent.session.id" = %session_id,
        );
        async move {
            let session_lock = self.session_lock(&session_id).await;
            let _guard = session_lock.lock().await;
            self.run_unlocked(agent_state, cancellation_token).await
        }
        .instrument(span)
        .await
    }

    async fn run_unlocked(
        &self,
        agent_state: Agent,
        cancellation_token: CancellationToken,
    ) -> Result<Agent, AgentError> {
        let outcome = self.run_inner(agent_state, cancellation_token).await?;
        if matches!(outcome, Agent::Success(_) | Agent::Fail(_, _)) {
            let session_id = self.get_session_id(&outcome);
            if let Err(e) = self.checkpoint_storage.delete_checkpoint(&session_id).await {
                return Ok(storage_failure(AgentState::new(Some(session_id)), e));
            }
        }
        Ok(outcome)
    }

    /// 恢复被中断的会话:有 checkpoint 就从中断点接着跑,没有则什么都不做返回 `Ok(None)`。
    ///
    /// 这是与 [`submit`](Self::submit) 互补的入口 —— `resume` 续上**被中断的 in-flight 状态**,
    /// `submit` 表示**用户发了新消息**(会丢弃尚未恢复的中断态)。
    pub async fn resume(
        &self,
        session_id: &str,
        cancellation_token: CancellationToken,
    ) -> Result<Option<Agent>, AgentError> {
        let span = tracing::info_span!(
            "agent.resume",
            "gen_ai.conversation.id" = %session_id,
            "tiny_agent.session.id" = %session_id,
        );
        async move {
            let session_lock = self.session_lock(session_id).await;
            let _guard = session_lock.lock().await;
            let checkpoint = match self.checkpoint_storage.get_checkpoint(session_id).await {
                Ok(Some(checkpoint)) => checkpoint,
                Ok(None) => return Ok(None),
                Err(e) => {
                    return Ok(Some(storage_failure(
                        AgentState::new(Some(session_id.to_string())),
                        e,
                    )));
                }
            };
            self.run_unlocked(checkpoint, cancellation_token)
                .await
                .map(Some)
        }
        .instrument(span)
        .await
    }

    /// 一站式入口:有 checkpoint 则恢复,否则从全新 `Ready` 开始,跑到停车点。
    /// 适合"按 session_id 拉起一轮"的调用方,无需自己拼装 `Agent`。
    pub async fn run_session(
        &self,
        session_id: &str,
        cancellation_token: CancellationToken,
    ) -> Result<Agent, AgentError> {
        let span = tracing::info_span!(
            "agent.run_session",
            "gen_ai.conversation.id" = %session_id,
            "tiny_agent.session.id" = %session_id,
        );
        async move {
            let session_lock = self.session_lock(session_id).await;
            let _guard = session_lock.lock().await;
            let start = match self.checkpoint_storage.get_checkpoint(session_id).await {
                Ok(Some(checkpoint)) => checkpoint,
                Ok(None) => Agent::Ready(AgentState::new(Some(session_id.to_string()))),
                Err(e) => {
                    return Ok(storage_failure(
                        AgentState::new(Some(session_id.to_string())),
                        e,
                    ));
                }
            };
            self.run_unlocked(start, cancellation_token).await
        }
        .instrument(span)
        .await
    }

    async fn run_inner(
        &self,
        mut agent_state: Agent,
        cancellation_token: CancellationToken,
    ) -> Result<Agent, AgentError> {
        loop {
            // ToolExecuting 仍有未应答的 tool_use,取消时不能在这里直接丢弃 ——
            // 必须让它流进 execute_pending_tools,由那里为每个 pending call 补占位结果,
            // 否则 transcript 会留下孤儿 tool_use。其余状态没有这个负担,可就地停车。
            if cancellation_token.is_cancelled()
                && !matches!(agent_state, Agent::ToolExecuting(_, _))
            {
                // 已经是 Interrupted 的状态:产生它的那一步(execute_pending_tools 里的
                // `interrupt()` / `save_interruption`)已经存档并 emit 过 Interrupted 了。
                // 这里若再 park、再 emit 一遍,取消恰好落在工具执行中时就会冒出**两条**
                // Interrupted 事件。已停车的状态直接原样返回即可。
                if matches!(agent_state, Agent::Interrupted(_)) {
                    return Ok(agent_state);
                }
                let session_id = self.get_session_id(&agent_state);
                let parked = match agent_state {
                    Agent::Reasoning(ctx) => Agent::Interrupted(ctx),
                    other => other,
                };
                if let Err(e) = self
                    .checkpoint_storage
                    .as_ref()
                    .save_checkpoint(&session_id, &parked)
                    .await
                {
                    return Ok(storage_failure(AgentState::new(Some(session_id)), e));
                }
                observability::interrupted(&session_id);
                return Ok(parked);
            }

            agent_state = match agent_state {
                Agent::Interrupted(ctx) => Agent::Ready(ctx),
                Agent::Ready(mut ctx) => {
                    // 迭代上限:每进入一次推理算一步。超限直接终态失败(不重试)
                    if ctx.steps_count >= self.max_iter {
                        let failure = AgentFailure::new(
                            FailureKind::MaxIter,
                            format!("exceeded max_iter ({})", self.max_iter),
                        );
                        observability::failed(
                            &ctx.session_id,
                            failure.kind.as_str(),
                            &failure.message,
                        );
                        return Ok(Agent::Fail(ctx, failure));
                    }
                    ctx.steps_count += 1;
                    Agent::Reasoning(ctx)
                }
                Agent::Reasoning(ctx) => {
                    self.handle_reasoning(ctx, cancellation_token.clone())
                        .await?
                }
                Agent::ToolExecuting(ctx, tool_calls) => {
                    self.execute_pending_tools(&ctx, tool_calls, cancellation_token.clone())
                        .await?
                }
                Agent::Fail(ctx, failure) => match self.handle_fail(ctx, failure).await? {
                    ready @ Agent::Ready(_) => ready,
                    terminal => return Ok(terminal),
                },
                other => return Ok(other),
            }
        }
    }

    pub async fn create_session(&self) -> Result<String, AgentError> {
        let new_session_id = Uuid::new_v4().to_string();
        let transcript = Transcript::new(&new_session_id);
        self.transcript_storage
            .save_transcript(&new_session_id, &transcript)
            .await?;
        let sandbox = self.sandbox.clone().open(&new_session_id).await?;
        sandbox.save().await?;
        Ok(new_session_id)
    }

    /// 框架对外的"用户说话"入口:把一条用户消息追加进会话历史,
    /// 返回一个 `Ready` 状态交给 `run` 驱动一轮对话。
    pub async fn submit(
        &self,
        session_id: &str,
        user_text: impl Into<String>,
    ) -> Result<Agent, AgentError> {
        let user_text = user_text.into();
        let span = tracing::info_span!(
            "user_input",
            "gen_ai.operation.name" = "user_input",
            "gen_ai.conversation.id" = %session_id,
            // 本轮用户原始输入；完整会话历史在随后的 gen_ai.chat span 的 input.messages 上。
            "gen_ai.input.messages" = %user_text,
            "tiny_agent.session.id" = %session_id,
        );
        self.submit_inner(session_id, user_text)
            .instrument(span)
            .await
    }

    async fn submit_inner(&self, session_id: &str, user_text: String) -> Result<Agent, AgentError> {
        let session_lock = self.session_lock(session_id).await;
        let _guard = session_lock.lock().await;
        if let Err(e) = self.checkpoint_storage.delete_checkpoint(session_id).await {
            return Ok(storage_failure(
                AgentState::new(Some(session_id.to_string())),
                e,
            ));
        }
        if let Err(e) = self
            .transcript_storage
            .append_message(session_id, Message::text(Role::User, user_text.clone()))
            .await
        {
            return Ok(storage_failure(
                AgentState::new(Some(session_id.to_string())),
                e,
            ));
        }
        self.messages.send(
            session_id,
            crate::messages::MessageKind::UserMessage {
                text: user_text.clone(),
            },
        );
        observability::user_message(session_id, &user_text);
        Ok(Agent::Ready(AgentState::new(Some(session_id.to_string()))))
    }

    async fn session_lock(&self, session_id: &str) -> std::sync::Arc<tokio::sync::Mutex<()>> {
        let mut locks = self.session_locks.lock().await;
        locks.retain(|_, lock| lock.strong_count() > 0);
        if let Some(lock) = locks.get(session_id).and_then(|lock| lock.upgrade()) {
            return lock;
        }
        let lock = std::sync::Arc::new(tokio::sync::Mutex::new(()));
        locks.insert(session_id.to_string(), std::sync::Arc::downgrade(&lock));
        lock
    }

    pub(crate) fn get_session_id(&self, agent: &Agent) -> String {
        match agent {
            Agent::Fail(ctx, _)
            | Agent::Interrupted(ctx)
            | Agent::Ready(ctx)
            | Agent::Reasoning(ctx)
            | Agent::Success(ctx)
            | Agent::ToolExecuting(ctx, _)
            | Agent::WaitingForUser(ctx, _) => ctx.session_id.clone(),
        }
    }

    /// 存档 checkpoint 并返回 `Interrupted`。调用前须确保所有 tool_use 已被应答,
    /// 因此 checkpoint 不再携带可重跑的 pending calls,resume 时只会从 `Interrupted`
    /// 走 `Ready → Reasoning`,不会重复执行工具。
    pub(crate) async fn interrupt(&self, ctx: &AgentState) -> Result<Agent, AgentError> {
        let agent = Agent::Interrupted(ctx.clone());
        if let Err(e) = self
            .checkpoint_storage
            .as_ref()
            .save_checkpoint(&ctx.session_id, &agent)
            .await
        {
            return Ok(storage_failure(ctx.clone(), e));
        }
        observability::interrupted(&ctx.session_id);
        Ok(agent)
    }

    pub(crate) async fn save_interruption(&self, agent: &Agent) -> Result<(), AgentError> {
        let session_id = self.get_session_id(agent);
        self.checkpoint_storage
            .as_ref()
            .save_checkpoint(&session_id, agent)
            .await?;
        observability::interrupted(&session_id);
        Ok(())
    }
}