tiny-agent 0.2.0

mod default_tools;

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

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

/// 调用工具时注入的运行时上下文。
///
/// 大多数工具用不到它(用 [`ToolRegistry::register`] 注册的 handler 拿不到);
/// 需要往前端推实时消息、或管理后台任务的工具(如 `bash`/`check`/`kill`)用
/// [`ToolRegistry::register_with_ctx`] 注册,才会收到这个上下文。
#[derive(Clone, Default)]
pub struct ToolCtx {
    /// 对前端的实时消息通道。
    pub messages: MessageChannel,
    /// 当前会话 id。
    pub session_id: String,
    /// 后台任务登记表(跨工具调用共享)。
    pub tasks: TaskManager,
}

type BoxFuture = Pin<Box<dyn Future<Output = Result<ToolOutcome, String>> + Send>>;
type BoxedCall = Arc<dyn Fn(Arc<dyn Sandbox>, Value, ToolCtx) -> 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, _ctx: ToolCtx| {
                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,
            },
        );
    }

    /// 与 [`register`](Self::register) 相同,但 handler 额外收到一个 [`ToolCtx`]
    /// (消息通道 / 会话 id / 后台任务表)。供 `bash`/`check`/`kill` 这类需要实时
    /// 推送或管理后台进程的工具使用。
    pub fn register_with_ctx<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, ToolCtx) -> 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, ctx: ToolCtx| {
            let handler = handler.clone();
            Box::pin(async move {
                let args: A = serde_json::from_value(v).map_err(|e| format!("参数不合法: {e}"))?;
                handler(sandbox, args, ctx).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 与 [`ToolCtx`]。
    pub async fn call(
        &self,
        name: &str,
        args: Value,
        sandbox: Arc<dyn Sandbox>,
        ctx: ToolCtx,
    ) -> Result<ToolOutcome, String> {
        let tool = self
            .tools
            .get(name)
            .ok_or_else(|| format!("工具不存在: {name}"))?;
        (tool.call)(sandbox, args, ctx).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()
    }
}