tiny-agent 0.1.0

mod default_tools;

pub use default_tools::{register_default_tools, register_skill_tools, register_spawn_agent};

use super::sandbox::Sandbox;
use crate::shared::{ToolDefinition, UserInteraction};
use schemars::{JsonSchema, schema_for};
use serde::de::DeserializeOwned;
use serde_json::Value;
use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};

type BoxFuture = Pin<Box<dyn Future<Output = Result<ToolOutcome, String>> + Send>>;
type BoxedCall = Arc<dyn Fn(Arc<dyn Sandbox>, Value) -> BoxFuture + Send + Sync>;

#[derive(Debug, Clone)]
pub enum ToolOutcome {
    Completed(Value),
    NeedsUserInteraction(UserInteraction),
}

impl From<Value> for ToolOutcome {
    fn from(value: Value) -> Self {
        Self::Completed(value)
    }
}

#[derive(Clone)]
pub struct ToolEntry {
    pub name: String,
    pub description: String,
    pub schema: Value,
    /// 是否只读、无副作用。只读工具之间可以并发执行(见 `core::tool_execution` 的分段调度);
    /// 任何会改动 sandbox 状态的工具都必须保持 `false`,以免与并发读发生 read-write race。
    pub read_only: bool,
    call: BoxedCall,
}

#[derive(Clone)]
pub struct ToolRegistry {
    tools: HashMap<String, ToolEntry>,
}

impl ToolRegistry {
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
        }
    }

    /// 注册一个工具。`read_only` 标记它是否无副作用 —— 只读工具之间可被并发调度
    /// (见 `core::tool_execution` 的分段执行),任何会改动 sandbox 的工具都必须传 `false`。
    /// handler 统一返回 `ToolOutcome`(普通完成用 `Value.into()`,需要用户交互用 `NeedsUserInteraction`)。
    pub fn register<A, F, Fut>(
        &mut self,
        name: &str,
        description: &str,
        read_only: bool,
        handler: F,
    ) where
        A: DeserializeOwned + JsonSchema + Send + 'static,
        F: Fn(Arc<dyn Sandbox>, A) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<ToolOutcome, String>> + Send + 'static,
    {
        let schema = serde_json::to_value(schema_for!(A)).unwrap();
        let handler = Arc::new(handler);
        let call: BoxedCall = Arc::new(move |sandbox: Arc<dyn Sandbox>, v: Value| {
            let handler = handler.clone();
            Box::pin(async move {
                let args: A = serde_json::from_value(v).map_err(|e| format!("参数不合法: {e}"))?;
                handler(sandbox, args).await
            })
        });
        self.tools.insert(
            name.into(),
            ToolEntry {
                name: name.into(),
                description: description.into(),
                schema,
                read_only,
                call,
            },
        );
    }

    /// 工具是否声明为只读。未知工具按非只读处理(最保守,强制串行)。
    pub fn is_read_only(&self, name: &str) -> bool {
        self.tools.get(name).is_some_and(|t| t.read_only)
    }

    /// 按名字调度执行,执行时注入该会话绑定的 sandbox。
    pub async fn call(
        &self,
        name: &str,
        args: Value,
        sandbox: Arc<dyn Sandbox>,
    ) -> Result<ToolOutcome, String> {
        let tool = self
            .tools
            .get(name)
            .ok_or_else(|| format!("工具不存在: {name}"))?;
        (tool.call)(sandbox, args).await
    }

    /// 将工具 struct 转成 `ToolDefinition` 列表
    pub fn definitions(&self) -> Vec<ToolDefinition> {
        self.tools
            .values()
            .map(|t| ToolDefinition {
                name: t.name.clone(),
                desc: t.description.clone(),
                arguments: t.schema.clone(),
            })
            .collect()
    }
}

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