tiny-agent 0.1.0

use super::{Agent, AgentRunTime};
use crate::{
    error::AgentError,
    shared::{Message, Role, ToolCall},
    state::AgentState,
    trajectory::TrajectoryEventKind,
};
use serde::{Deserialize, Serialize};
use std::fmt;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

/// 退避时长上限,避免指数爆炸把单次重试拖到分钟级。
const MAX_BACKOFF: Duration = Duration::from_secs(30);

/// 失败的结构化类别。
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FailureKind {
    /// 达到最大迭代步数。
    MaxIter,
    /// 调用 provider 建立流失败(网络 / 传输层)。
    Provider,
    /// provider 明确返回确定性失败,例如鉴权失败、请求格式错误、模型不存在。
    ProviderPermanent,
    /// provider 限流。
    ProviderRateLimited,
    /// 流式响应过程中出错,或未收到正常收尾。
    ProviderStream,
    /// 模型输出被长度上限截断。
    Truncated,
    /// 模型输出被内容过滤拦截。
    ContentFilter,
    /// 中间件自身报错。
    Middleware,
    /// 工具执行失败(模型可据错误结果调整后继续)。
    Tool,
    /// 打开沙箱失败。
    SandboxOpen,
    /// 保存沙箱状态失败。
    SandboxSave,
    /// transcript/checkpoint 等持久化层失败。
    Storage,
}

impl FailureKind {
    /// 该类失败是否值得用同样的输入重试。
    ///
    /// 可重试:瞬时性的 provider / 流 / 沙箱故障,以及工具失败(交回模型自愈)。
    /// 不可重试:`MaxIter` / `Truncated` / `ContentFilter` —— 重试必然复现,
    /// 应直接进终态。`Middleware` 视为确定性错误(护栏 / 逻辑 bug),同样终态。
    pub fn is_retryable(&self) -> bool {
        matches!(
            self,
            FailureKind::Provider
                | FailureKind::ProviderRateLimited
                | FailureKind::ProviderStream
                | FailureKind::Tool
                | FailureKind::SandboxOpen
                | FailureKind::SandboxSave
                | FailureKind::Storage
        )
    }

    /// 稳定的 snake_case 字符串,用于 trajectory 事件等线上格式。
    pub fn as_str(&self) -> &'static str {
        match self {
            FailureKind::MaxIter => "max_iter",
            FailureKind::Provider => "provider",
            FailureKind::ProviderPermanent => "provider_permanent",
            FailureKind::ProviderRateLimited => "provider_rate_limited",
            FailureKind::ProviderStream => "provider_stream",
            FailureKind::Truncated => "truncated",
            FailureKind::ContentFilter => "content_filter",
            FailureKind::Middleware => "middleware",
            FailureKind::Tool => "tool",
            FailureKind::SandboxOpen => "sandbox_open",
            FailureKind::SandboxSave => "sandbox_save",
            FailureKind::Storage => "storage",
        }
    }
}

impl fmt::Display for FailureKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

#[derive(thiserror::Error, Debug, Serialize, Deserialize, Clone)]
#[error("{kind}: {message}")]
pub struct AgentFailure {
    pub kind: FailureKind,
    pub message: String,
    pub tool_call: Option<ToolCall>,
}

impl AgentFailure {
    pub(crate) fn new(kind: FailureKind, message: impl Into<String>) -> Self {
        Self {
            kind,
            message: message.into(),
            tool_call: None,
        }
    }

    pub(crate) fn tool(tool_call: ToolCall, message: impl Into<String>) -> Self {
        Self {
            kind: FailureKind::Tool,
            message: message.into(),
            tool_call: Some(tool_call),
        }
    }
}

impl AgentRunTime {
    pub(crate) fn fail_with(
        &self,
        ctx: AgentState,
        kind: FailureKind,
        message: impl Into<String>,
    ) -> Agent {
        Agent::Fail(ctx, AgentFailure::new(kind, message))
    }

    pub(crate) fn fail_tool(
        &self,
        ctx: AgentState,
        tool_call: ToolCall,
        message: impl Into<String>,
    ) -> Agent {
        Agent::Fail(ctx, AgentFailure::tool(tool_call, message))
    }

    pub(crate) async fn handle_fail(
        &self,
        mut ctx: AgentState,
        mut failure: AgentFailure,
    ) -> Result<Agent, AgentError> {
        loop {
            self.emit(
                &ctx.session_id,
                TrajectoryEventKind::Failed {
                    kind: failure.kind.as_str().to_string(),
                    message: failure.message.clone(),
                },
            )
            .await;

            // 不可重试的类别:重试只会用相同输入复现,直接进终态,不消耗预算、不退避。
            if !failure.kind.is_retryable() {
                return Ok(Agent::Fail(ctx, failure));
            }

            ctx.consecutive_fail_count = ctx.consecutive_fail_count.saturating_add(1);
            if ctx.consecutive_fail_count > self.max_consecutive_fail_count {
                return Ok(Agent::Fail(ctx, failure));
            }

            // 工具失败的 `tool_result` 已由 execute_pending_tools 写入 transcript
            // (那里负责维持"每个 tool_use 必有 tool_result"的不变量),这里不再重复回应。
            // 非工具失败(provider / stream / sandbox 等)没有对应的 tool_use,
            // 补一条 system 面包屑让模型知道发生了什么并据此重试。
            //
            // Storage 失败时不要再写 transcript:持久化层已经在抖,继续写只会把
            // 可重试失败变成外层 abort。
            if failure.tool_call.is_none() && failure.kind != FailureKind::Storage {
                if let Err(e) = self
                    .transcript_storage
                    .append_message(
                        &ctx.session_id,
                        Message::text(
                            Role::System,
                            format!("{} failed: {}", failure.kind, failure.message),
                        ),
                    )
                    .await
                {
                    failure = AgentFailure::new(FailureKind::Storage, e.to_string());
                    continue;
                }
            }

            // 重试前退避:指数 + 等量 jitter,打散同时失败的并发会话。
            self.retry_backoff(ctx.consecutive_fail_count).await;
            return Ok(Agent::Ready(ctx));
        }
    }

    /// 第 `attempt` 次连续失败后的退避等待:`base * 2^(attempt-1)`,封顶 `MAX_BACKOFF`,
    /// 再叠加 [delay/2, delay] 区间的 jitter。`base` 为 0 时直接返回(测试用)。
    async fn retry_backoff(&self, attempt: u32) {
        let base = self.retry_backoff;
        if base.is_zero() {
            return;
        }
        let exponent = attempt.saturating_sub(1).min(6); // 2^6 后由 MAX_BACKOFF 封顶
        let delay = base.saturating_mul(1u32 << exponent).min(MAX_BACKOFF);
        let half = delay / 2;
        let half_nanos = (half.as_nanos() as u64).max(1);
        let jitter_source = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.subsec_nanos() as u64)
            .unwrap_or(0);
        let jitter = Duration::from_nanos(jitter_source % half_nanos);
        tokio::time::sleep(half + jitter).await;
    }
}