tiny-agent 0.1.0

use super::{AgentRunTime, prompt::PromptComposer};
use crate::{
    checkpoint::CheckpointStorage,
    error::AgentError,
    middleware::{Middleware, MiddlewareChain},
    providers::LlmProvider,
    sandbox::Sandbox,
    skills::SkillManager,
    tools::{ToolRegistry, register_skill_tools},
    trajectory::{NoopSink, TrajectorySink},
    transcript::TranscriptStorage,
};
use std::{collections::HashMap, path::PathBuf, sync::Arc, time::Duration};
use tokio::sync::Mutex;

const DEFAULT_MAX_ITER: u32 = 10;
const DEFAULT_MAX_CONSECUTIVE_FAIL_COUNT: u32 = 3;
const DEFAULT_RETRY_BACKOFF: Duration = Duration::from_millis(200);
const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(120);

#[derive(Clone)]
pub struct AgentRuntimeConfig {
    pub provider: Arc<dyn LlmProvider>,
    pub model: String,
    pub system_prompt: Option<String>,
    pub temperature: Option<f32>,
    pub max_tokens: Option<u32>,
    pub max_iter: u32,
    pub max_consecutive_fail_count: u32,
    pub checkpoint_storage: Arc<dyn CheckpointStorage>,
    pub transcript_storage: Arc<dyn TranscriptStorage>,
    pub sandbox: Arc<dyn Sandbox>,
    pub tools: ToolRegistry,
    pub skill_dirs: Vec<PathBuf>,
    pub middleware: MiddlewareChain,
    pub trajectory: Arc<dyn TrajectorySink>,
    pub retry_backoff: Duration,
    pub tool_timeout: Duration,
}

impl AgentRuntimeConfig {
    pub fn new(
        provider: Arc<dyn LlmProvider>,
        model: impl Into<String>,
        checkpoint_storage: Arc<dyn CheckpointStorage>,
        transcript_storage: Arc<dyn TranscriptStorage>,
        sandbox: Arc<dyn Sandbox>,
    ) -> Self {
        Self {
            provider,
            model: model.into(),
            system_prompt: None,
            temperature: None,
            max_tokens: None,
            max_iter: DEFAULT_MAX_ITER,
            max_consecutive_fail_count: DEFAULT_MAX_CONSECUTIVE_FAIL_COUNT,
            checkpoint_storage,
            transcript_storage,
            sandbox,
            tools: ToolRegistry::new(),
            skill_dirs: Vec::new(),
            middleware: MiddlewareChain::new(),
            trajectory: Arc::new(NoopSink),
            retry_backoff: DEFAULT_RETRY_BACKOFF,
            tool_timeout: DEFAULT_TOOL_TIMEOUT,
        }
    }

    pub fn with_tools(mut self, tools: ToolRegistry) -> Self {
        self.tools = tools;
        self
    }

    pub fn with_system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
        self.system_prompt = Some(system_prompt.into());
        self
    }

    pub fn with_middleware(mut self, middleware: Arc<dyn Middleware>) -> Self {
        self.middleware.push(middleware);
        self
    }

    pub fn with_trajectory(mut self, trajectory: Arc<dyn TrajectorySink>) -> Self {
        self.trajectory = trajectory;
        self
    }
}

impl AgentRunTime {
    pub fn from_config(config: AgentRuntimeConfig) -> Result<Self, AgentError> {
        let mut tools = config.tools;
        let skills = if config.skill_dirs.is_empty() {
            None
        } else {
            Some(Arc::new(SkillManager::discover_dirs(&config.skill_dirs)?))
        };
        if let Some(skills) = &skills {
            register_skill_tools(&mut tools, skills.clone());
        }

        Ok(Self {
            provider: config.provider,
            max_iter: config.max_iter,
            max_consecutive_fail_count: config.max_consecutive_fail_count,
            model: config.model,
            system_prompt: config.system_prompt,
            temperature: config.temperature,
            max_tokens: config.max_tokens,
            checkpoint_storage: config.checkpoint_storage,
            transcript_storage: config.transcript_storage,
            sandbox: config.sandbox,
            tools,
            skills,
            prompt_composer: PromptComposer,
            middleware: config.middleware,
            trajectory: config.trajectory,
            retry_backoff: config.retry_backoff,
            tool_timeout: config.tool_timeout,
            session_locks: Mutex::new(HashMap::new()),
        })
    }
}

/// `AgentRunTime` 的构造器。
///
/// `provider` / `model` / 三个存储 / `sandbox` 是必填项;
/// 中间件与轨迹 sink 可选(默认空链 + `NoopSink`)。
///
/// ```ignore
/// let runtime = AgentRunTimeBuilder::new()
///     .provider(provider)
///     .model("gpt-4o")
///     .checkpoint_storage(storage.clone())
///     .transcript_storage(storage)
///     .sandbox(sandbox)
///     .tools(tools)
///     .middleware(Arc::new(MyGuardrail))
///     .trajectory(sink)
///     .build()?;
/// ```
pub struct AgentRunTimeBuilder {
    provider: Option<Arc<dyn LlmProvider>>,
    model: Option<String>,
    system_prompt: Option<String>,
    temperature: Option<f32>,
    max_tokens: Option<u32>,
    max_iter: u32,
    max_consecutive_fail_count: u32,
    checkpoint_storage: Option<Arc<dyn CheckpointStorage>>,
    transcript_storage: Option<Arc<dyn TranscriptStorage>>,
    sandbox: Option<Arc<dyn Sandbox>>,
    tools: ToolRegistry,
    skill_dirs: Vec<PathBuf>,
    middleware: MiddlewareChain,
    trajectory: Option<Arc<dyn TrajectorySink>>,
    retry_backoff: Duration,
    tool_timeout: Duration,
}

impl AgentRunTimeBuilder {
    pub fn new() -> Self {
        Self {
            provider: None,
            model: None,
            system_prompt: None,
            temperature: None,
            max_tokens: None,
            max_iter: DEFAULT_MAX_ITER,
            max_consecutive_fail_count: DEFAULT_MAX_CONSECUTIVE_FAIL_COUNT,
            checkpoint_storage: None,
            transcript_storage: None,
            sandbox: None,
            tools: ToolRegistry::new(),
            skill_dirs: Vec::new(),
            middleware: MiddlewareChain::new(),
            trajectory: None,
            retry_backoff: DEFAULT_RETRY_BACKOFF,
            tool_timeout: DEFAULT_TOOL_TIMEOUT,
        }
    }

    pub fn provider(mut self, provider: Arc<dyn LlmProvider>) -> Self {
        self.provider = Some(provider);
        self
    }

    pub fn model(mut self, model: impl Into<String>) -> Self {
        self.model = Some(model.into());
        self
    }

    pub fn system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
        self.system_prompt = Some(system_prompt.into());
        self
    }

    pub fn with_system_prompt(self, system_prompt: impl Into<String>) -> Self {
        self.system_prompt(system_prompt)
    }

    pub fn temperature(mut self, temperature: f32) -> Self {
        self.temperature = Some(temperature);
        self
    }

    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = Some(max_tokens);
        self
    }

    pub fn max_iter(mut self, max_iter: u32) -> Self {
        self.max_iter = max_iter;
        self
    }

    pub fn max_consecutive_fail_count(mut self, count: u32) -> Self {
        self.max_consecutive_fail_count = count;
        self
    }

    pub fn checkpoint_storage(mut self, storage: Arc<dyn CheckpointStorage>) -> Self {
        self.checkpoint_storage = Some(storage);
        self
    }

    pub fn transcript_storage(mut self, storage: Arc<dyn TranscriptStorage>) -> Self {
        self.transcript_storage = Some(storage);
        self
    }

    pub fn sandbox(mut self, sandbox: Arc<dyn Sandbox>) -> Self {
        self.sandbox = Some(sandbox);
        self
    }

    pub fn tools(mut self, tools: ToolRegistry) -> Self {
        self.tools = tools;
        self
    }

    pub fn skill_dir(mut self, path: impl Into<PathBuf>) -> Self {
        self.skill_dirs.push(path.into());
        self
    }

    pub fn skill_dirs<I, P>(mut self, paths: I) -> Self
    where
        I: IntoIterator<Item = P>,
        P: Into<PathBuf>,
    {
        self.skill_dirs.extend(paths.into_iter().map(Into::into));
        self
    }

    /// 追加一个中间件(可多次调用,按调用顺序构成 onion 链)。
    pub fn middleware(mut self, middleware: Arc<dyn Middleware>) -> Self {
        self.middleware.push(middleware);
        self
    }

    /// 设置轨迹 sink(不设则为 `NoopSink`)。
    pub fn trajectory(mut self, trajectory: Arc<dyn TrajectorySink>) -> Self {
        self.trajectory = Some(trajectory);
        self
    }

    /// 可重试失败的退避基数(默认 200ms,实际等待为指数退避 + jitter)。
    /// 传 `Duration::ZERO` 可关闭退避(测试常用)。
    pub fn retry_backoff(mut self, base: Duration) -> Self {
        self.retry_backoff = base;
        self
    }

    /// 单个工具调用的超时时间(默认 120s)。传 `Duration::ZERO` 可关闭工具超时。
    pub fn tool_timeout(mut self, timeout: Duration) -> Self {
        self.tool_timeout = timeout;
        self
    }

    pub fn build(self) -> Result<AgentRunTime, AgentError> {
        AgentRunTime::from_config(AgentRuntimeConfig {
            provider: self
                .provider
                .ok_or_else(|| AgentError::config("provider is required"))?,
            model: self
                .model
                .ok_or_else(|| AgentError::config("model is required"))?,
            system_prompt: self.system_prompt,
            temperature: self.temperature,
            max_tokens: self.max_tokens,
            max_iter: self.max_iter,
            max_consecutive_fail_count: self.max_consecutive_fail_count,
            checkpoint_storage: self
                .checkpoint_storage
                .ok_or_else(|| AgentError::config("checkpoint_storage is required"))?,
            transcript_storage: self
                .transcript_storage
                .ok_or_else(|| AgentError::config("transcript_storage is required"))?,
            sandbox: self
                .sandbox
                .ok_or_else(|| AgentError::config("sandbox is required"))?,
            tools: self.tools,
            skill_dirs: self.skill_dirs,
            middleware: self.middleware,
            trajectory: self.trajectory.unwrap_or_else(|| Arc::new(NoopSink)),
            retry_backoff: self.retry_backoff,
            tool_timeout: self.tool_timeout,
        })
    }
}

impl Default for AgentRunTimeBuilder {
    fn default() -> Self {
        Self::new()
    }
}