tiny_agent/tools/
mod.rs

1mod default_tools;
2
3pub use default_tools::{register_default_tools, register_skill_tools, register_spawn_agent};
4
5use super::sandbox::Sandbox;
6use crate::messages::MessageChannel;
7use crate::shared::{ToolDefinition, UserInteraction};
8use crate::tasks::TaskManager;
9use schemars::{JsonSchema, schema_for};
10use serde::de::DeserializeOwned;
11use serde_json::Value;
12use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};
13
14/// 调用工具时注入的运行时上下文。
15///
16/// 大多数工具用不到它(用 [`ToolRegistry::register`] 注册的 handler 拿不到);
17/// 需要往前端推实时消息、或管理后台任务的工具(如 `bash`/`check`/`kill`)用
18/// [`ToolRegistry::register_with_ctx`] 注册,才会收到这个上下文。
19#[derive(Clone, Default)]
20pub struct ToolCtx {
21    /// 对前端的实时消息通道。
22    pub messages: MessageChannel,
23    /// 当前会话 id。
24    pub session_id: String,
25    /// 后台任务登记表(跨工具调用共享)。
26    pub tasks: TaskManager,
27}
28
29type BoxFuture = Pin<Box<dyn Future<Output = Result<ToolOutcome, String>> + Send>>;
30type BoxedCall = Arc<dyn Fn(Arc<dyn Sandbox>, Value, ToolCtx) -> BoxFuture + Send + Sync>;
31
32#[derive(Debug, Clone)]
33pub enum ToolOutcome {
34    Completed(Value),
35    NeedsUserInteraction(UserInteraction),
36}
37
38impl From<Value> for ToolOutcome {
39    fn from(value: Value) -> Self {
40        Self::Completed(value)
41    }
42}
43
44#[derive(Clone)]
45pub struct ToolEntry {
46    pub name: String,
47    pub description: String,
48    pub schema: Value,
49    /// 是否只读、无副作用。只读工具之间可以并发执行(见 `core::tool_execution` 的分段调度);
50    /// 任何会改动 sandbox 状态的工具都必须保持 `false`,以免与并发读发生 read-write race。
51    pub read_only: bool,
52    call: BoxedCall,
53}
54
55#[derive(Clone)]
56pub struct ToolRegistry {
57    tools: HashMap<String, ToolEntry>,
58}
59
60impl ToolRegistry {
61    pub fn new() -> Self {
62        Self {
63            tools: HashMap::new(),
64        }
65    }
66
67    /// 注册一个工具。`read_only` 标记它是否无副作用 —— 只读工具之间可被并发调度
68    /// (见 `core::tool_execution` 的分段执行),任何会改动 sandbox 的工具都必须传 `false`。
69    /// handler 统一返回 `ToolOutcome`(普通完成用 `Value.into()`,需要用户交互用 `NeedsUserInteraction`)。
70    pub fn register<A, F, Fut>(
71        &mut self,
72        name: &str,
73        description: &str,
74        read_only: bool,
75        handler: F,
76    ) where
77        A: DeserializeOwned + JsonSchema + Send + 'static,
78        F: Fn(Arc<dyn Sandbox>, A) -> Fut + Send + Sync + 'static,
79        Fut: Future<Output = Result<ToolOutcome, String>> + Send + 'static,
80    {
81        let schema = serde_json::to_value(schema_for!(A)).unwrap();
82        let handler = Arc::new(handler);
83        let call: BoxedCall =
84            Arc::new(move |sandbox: Arc<dyn Sandbox>, v: Value, _ctx: ToolCtx| {
85                let handler = handler.clone();
86                Box::pin(async move {
87                    let args: A =
88                        serde_json::from_value(v).map_err(|e| format!("参数不合法: {e}"))?;
89                    handler(sandbox, args).await
90                })
91            });
92        self.tools.insert(
93            name.into(),
94            ToolEntry {
95                name: name.into(),
96                description: description.into(),
97                schema,
98                read_only,
99                call,
100            },
101        );
102    }
103
104    /// 与 [`register`](Self::register) 相同,但 handler 额外收到一个 [`ToolCtx`]
105    /// (消息通道 / 会话 id / 后台任务表)。供 `bash`/`check`/`kill` 这类需要实时
106    /// 推送或管理后台进程的工具使用。
107    pub fn register_with_ctx<A, F, Fut>(
108        &mut self,
109        name: &str,
110        description: &str,
111        read_only: bool,
112        handler: F,
113    ) where
114        A: DeserializeOwned + JsonSchema + Send + 'static,
115        F: Fn(Arc<dyn Sandbox>, A, ToolCtx) -> Fut + Send + Sync + 'static,
116        Fut: Future<Output = Result<ToolOutcome, String>> + Send + 'static,
117    {
118        let schema = serde_json::to_value(schema_for!(A)).unwrap();
119        let handler = Arc::new(handler);
120        let call: BoxedCall = Arc::new(move |sandbox: Arc<dyn Sandbox>, v: Value, ctx: ToolCtx| {
121            let handler = handler.clone();
122            Box::pin(async move {
123                let args: A = serde_json::from_value(v).map_err(|e| format!("参数不合法: {e}"))?;
124                handler(sandbox, args, ctx).await
125            })
126        });
127        self.tools.insert(
128            name.into(),
129            ToolEntry {
130                name: name.into(),
131                description: description.into(),
132                schema,
133                read_only,
134                call,
135            },
136        );
137    }
138
139    /// 工具是否声明为只读。未知工具按非只读处理(最保守,强制串行)。
140    pub fn is_read_only(&self, name: &str) -> bool {
141        self.tools.get(name).is_some_and(|t| t.read_only)
142    }
143
144    /// 按名字调度执行,执行时注入该会话绑定的 sandbox 与 [`ToolCtx`]。
145    pub async fn call(
146        &self,
147        name: &str,
148        args: Value,
149        sandbox: Arc<dyn Sandbox>,
150        ctx: ToolCtx,
151    ) -> Result<ToolOutcome, String> {
152        let tool = self
153            .tools
154            .get(name)
155            .ok_or_else(|| format!("工具不存在: {name}"))?;
156        (tool.call)(sandbox, args, ctx).await
157    }
158
159    /// 将工具 struct 转成 `ToolDefinition` 列表
160    pub fn definitions(&self) -> Vec<ToolDefinition> {
161        self.tools
162            .values()
163            .map(|t| ToolDefinition {
164                name: t.name.clone(),
165                desc: t.description.clone(),
166                arguments: t.schema.clone(),
167            })
168            .collect()
169    }
170}
171
172impl Default for ToolRegistry {
173    fn default() -> Self {
174        Self::new()
175    }
176}
tiny_agent/tools/mod.rs

tiny_agent/tools/
mod.rs
