tiny-agent 0.2.0

use super::{AgentRunTime, prompt::PromptComposer};
use crate::{
    checkpoint::CheckpointStorage,
    error::AgentError,
    messages::MessageChannel,
    middleware::{Middleware, MiddlewareChain},
    providers::LlmProvider,
    sandbox::Sandbox,
    skills::SkillManager,
    tools::{ToolRegistry, register_skill_tools},
    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 messages: MessageChannel,
    pub retry_backoff: Duration,
    pub tool_timeout: Duration,
}

#[derive(Clone)]
pub struct AgentRuntimeStorage {
    pub checkpoint: Arc<dyn CheckpointStorage>,
    pub transcript: Arc<dyn TranscriptStorage>,
}

impl AgentRuntimeStorage {
    pub fn new(
        checkpoint: Arc<dyn CheckpointStorage>,
        transcript: Arc<dyn TranscriptStorage>,
    ) -> Self {
        Self {
            checkpoint,
            transcript,
        }
    }

    pub fn shared<S>(storage: Arc<S>) -> Self
    where
        S: CheckpointStorage + TranscriptStorage + 'static,
    {
        Self {
            checkpoint: storage.clone(),
            transcript: storage,
        }
    }
}

#[derive(Clone)]
pub struct AgentRuntimeOptions {
    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 tools: ToolRegistry,
    pub skill_dirs: Vec<PathBuf>,
    pub middleware: MiddlewareChain,
    pub messages: MessageChannel,
    pub retry_backoff: Duration,
    pub tool_timeout: Duration,
}

impl Default for AgentRuntimeOptions {
    fn default() -> Self {
        Self {
            system_prompt: None,
            temperature: None,
            max_tokens: None,
            max_iter: DEFAULT_MAX_ITER,
            max_consecutive_fail_count: DEFAULT_MAX_CONSECUTIVE_FAIL_COUNT,
            tools: ToolRegistry::new(),
            skill_dirs: Vec::new(),
            middleware: MiddlewareChain::new(),
            messages: MessageChannel::disconnected(),
            retry_backoff: DEFAULT_RETRY_BACKOFF,
            tool_timeout: DEFAULT_TOOL_TIMEOUT,
        }
    }
}

#[derive(Clone)]
pub struct AgentRuntimeConfigInput {
    pub provider: Arc<dyn LlmProvider>,
    pub model: String,
    pub storage: AgentRuntimeStorage,
    pub sandbox: Arc<dyn Sandbox>,
    pub options: AgentRuntimeOptions,
}

impl AgentRuntimeConfig {
    pub fn new(input: AgentRuntimeConfigInput) -> Self {
        Self {
            provider: input.provider,
            model: input.model,
            system_prompt: input.options.system_prompt,
            temperature: input.options.temperature,
            max_tokens: input.options.max_tokens,
            max_iter: input.options.max_iter,
            max_consecutive_fail_count: input.options.max_consecutive_fail_count,
            checkpoint_storage: input.storage.checkpoint,
            transcript_storage: input.storage.transcript,
            sandbox: input.sandbox,
            tools: input.options.tools,
            skill_dirs: input.options.skill_dirs,
            middleware: input.options.middleware,
            messages: input.options.messages,
            retry_backoff: input.options.retry_backoff,
            tool_timeout: input.options.tool_timeout,
        }
    }
}

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,
            messages: config.messages,
            tasks: crate::tasks::TaskManager::new(),
            retry_backoff: config.retry_backoff,
            tool_timeout: config.tool_timeout,
            session_locks: Mutex::new(HashMap::new()),
        })
    }
}

/// `AgentRunTime` 的构造器。
///
/// `provider` / `model` / 两类存储 / `sandbox` 是必填项;
/// 中间件与实时消息通道可选。可观测性由标准 `tracing` span/event 产出,
/// 是否导出到 OpenTelemetry / 日志 / 文件由宿主程序配置 subscriber 决定。
///
/// ```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))
///     .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,
    messages: Option<MessageChannel>,
    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(),
            messages: 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 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
    }

    /// 设置对前端的实时消息通道(不设则为 disconnected,消息静默丢弃)。
    pub fn messages(mut self, messages: MessageChannel) -> Self {
        self.messages = Some(messages);
        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::new(AgentRuntimeConfigInput {
            provider: self
                .provider
                .ok_or_else(|| AgentError::config("provider is required"))?,
            model: self
                .model
                .ok_or_else(|| AgentError::config("model is required"))?,
            storage: AgentRuntimeStorage::new(
                self.checkpoint_storage
                    .ok_or_else(|| AgentError::config("checkpoint_storage is required"))?,
                self.transcript_storage
                    .ok_or_else(|| AgentError::config("transcript_storage is required"))?,
            ),
            sandbox: self
                .sandbox
                .ok_or_else(|| AgentError::config("sandbox is required"))?,
            options: AgentRuntimeOptions {
                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,
                tools: self.tools,
                skill_dirs: self.skill_dirs,
                middleware: self.middleware,
                messages: self.messages.unwrap_or_default(),
                retry_backoff: self.retry_backoff,
                tool_timeout: self.tool_timeout,
            },
        }))
    }
}

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