AI Agent 工具注册表Rust trait 实现可插拔的工具发现机制一、AI Agent 的工具体系需要什么在 AI Agent 的架构里工具Tool是 Agent 和外部世界交互的接口。比如read_file—— 读取文件web_search—— 网络搜索run_code—— 执行代码query_database—— 查询数据库Agent 需要知道三件事有哪些工具、每个工具怎么描述给 AI 理解、怎么调用工具。flowchart TB subgraph Agent[AI Agent] LLM[大语言模型] Router[工具路由器] end subgraph Registry[工具注册表] T1[read_file] T2[web_search] T3[run_code] T4[query_database] T5[...] end LLM --|1. 返回函数调用| Router Router --|2. 查找工具| Registry Registry --|3. 返回工具实例| Router Router --|4. 执行工具| T3 T3 --|5. 返回结果| Router Router --|6. 结果反馈| LLM style Registry fill:#ffd93d,stroke:#333 style Router fill:#6bcb77,stroke:#333一个好的工具体系应该满足可发现Agent 能拿到所有工具的元信息名称、描述、参数 Schema可调用Agent 能通过统一接口调用任意工具可插拔新增/删除工具不需要改 Agent 核心代码类型安全工具的输入输出有明确类型编译期检查一个没有工具注册表的痛最早写 Agent 时工具发现是硬编码的。每加一个工具要改 Agent 核心里三处代码注册逻辑、描述生成、调用分发。加第三个工具时还不觉得烦加到第十个时已经在怀疑自己的设计。更难受的是别人想开发第三方工具必须 fork Agent 仓库不可能在运行时接一个新工具。这就是工具注册表的价值工具和 Agent 解耦。二、定义 Tool trait统一接口首先定义一个 trait所有工具都必须实现它use std::collections::HashMap; use serde::{Deserialize, Serialize}; use async_trait::async_trait; // 工具系统的核心类型定义 /// 工具参数的 JSON Schema —— 用于告诉 AI 这个工具需要什么参数 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolParamSchema { /// 参数名称 pub name: String, /// 参数类型: string | number | boolean | object | array #[serde(rename type)] pub param_type: String, /// 参数描述 pub description: String, /// 是否必须 pub required: bool, } /// 工具的元信息 —— 用于向 AI 描述这个工具 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolMetadata { /// 工具名称唯一标识 pub name: String, /// 工具描述 —— AI 根据这个决定是否调用 pub description: String, /// 参数列表 pub parameters: VecToolParamSchema, } /// 工具执行结果 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolResult { /// 是否成功 pub success: bool, /// 执行结果字符串表示方便 AI 理解 pub content: String, /// 附加的结构化数据可选 pub data: Optionserde_json::Value, /// 错误信息如果失败 pub error: OptionString, } /// 工具 trait —— 所有工具的统一接口 /// /// 使用 async_trait 宏是因为 Rust 原生 trait 不支持 async fn。 /// async_trait 会把 async fn 转换成返回 PinBoxdyn Future 的方法。 #[async_trait] pub trait Tool: Send Sync { /// 获取工具的元信息 fn metadata(self) - ToolMetadata; /// 执行工具 —— 接收 JSON 格式的参数返回执行结果 async fn execute(self, params: serde_json::Value) - ToolResult; /// 验证参数是否合法可选覆盖默认始终返回 Ok fn validate_params(self, _params: serde_json::Value) - Result(), String { Ok(()) } }三、实现几个具体工具工具 1文件读取工具// 具体工具实现 use std::path::PathBuf; /// 文件读取工具 pub struct ReadFileTool { /// 允许访问的根目录 —— 安全沙箱 allowed_root: PathBuf, } impl ReadFileTool { pub fn new(allowed_root: PathBuf) - Self { ReadFileTool { allowed_root } } /// 安全检查确保文件路径在允许的根目录内 fn canonicalize(self, path: str) - ResultPathBuf, String { let full_path self.allowed_root.join(path.trim_start_matches(/)); // canonicalize 会解析符号链接返回绝对路径 let canonical full_path .canonicalize() .map_err(|e| format!(路径解析失败: {e}))?; // 安全检查确保解析后的路径仍在 allowed_root 内 if !canonical.starts_with(self.allowed_root) { return Err(format!(路径越界: {} 不在 {}, canonical.display(), self.allowed_root.display())); } Ok(canonical) } } #[async_trait] impl Tool for ReadFileTool { fn metadata(self) - ToolMetadata { ToolMetadata { name: read_file.to_string(), description: 读取指定路径的文件内容.to_string(), parameters: vec![ ToolParamSchema { name: path.to_string(), param_type: string.to_string(), description: 要读取的文件路径.to_string(), required: true, }, ToolParamSchema { name: encoding.to_string(), param_type: string.to_string(), description: 文件编码默认 utf-8.to_string(), required: false, }, ], } } async fn execute(self, params: serde_json::Value) - ToolResult { // 提取 path 参数 let path match params.get(path).and_then(|v| v.as_str()) { Some(p) p, None return ToolResult { success: false, content: String::new(), data: None, error: Some(缺少必填参数: path.to_string()), }, }; // 安全检查 let file_path match self.canonicalize(path) { Ok(p) p, Err(e) return ToolResult { success: false, content: String::new(), data: None, error: Some(e), }, }; // 读取文件 match std::fs::read_to_string(file_path) { Ok(content) ToolResult { success: true, content: format!(文件 {} 内容:\n{}, path, content), data: Some(serde_json::json!({ path: path, size: content.len() })), error: None, }, Err(e) ToolResult { success: false, content: String::new(), data: None, error: Some(format!(读取文件失败: {e})), }, } } }工具 2网络搜索工具/// 网络搜索工具 pub struct WebSearchTool { /// API 密钥 api_key: String, /// HTTP 客户端 client: reqwest::Client, } impl WebSearchTool { pub fn new(api_key: String) - Self { WebSearchTool { api_key, client: reqwest::Client::new(), } } } #[async_trait] impl Tool for WebSearchTool { fn metadata(self) - ToolMetadata { ToolMetadata { name: web_search.to_string(), description: 搜索互联网获取最新信息.to_string(), parameters: vec![ ToolParamSchema { name: query.to_string(), param_type: string.to_string(), description: 搜索关键词.to_string(), required: true, }, ToolParamSchema { name: max_results.to_string(), param_type: number.to_string(), description: 最大结果数默认 5.to_string(), required: false, }, ], } } async fn execute(self, params: serde_json::Value) - ToolResult { let query match params.get(query).and_then(|v| v.as_str()) { Some(q) q, None return ToolResult { success: false, content: String::new(), data: None, error: Some(缺少必填参数: query.to_string()), }, }; let max params .get(max_results) .and_then(|v| v.as_u64()) .unwrap_or(5) as usize; // 模拟搜索实际应调用搜索 API let results vec![ format!(1. {} - 搜索结果摘要, query), format!(2. {} - 相关信息, query), ]; ToolResult { success: true, content: format!(搜索 \{}\ 的结果:\n{}, query, results.join(\n)), data: Some(serde_json::json!({ query: query, count: results.len() })), error: None, } } }工具 3代码执行工具沙箱内执行/// 代码执行工具 —— 在受控环境中执行代码 pub struct CodeRunnerTool { /// 执行超时秒 timeout: u64, } impl CodeRunnerTool { pub fn new(timeout: u64) - Self { CodeRunnerTool { timeout } } } #[async_trait] impl Tool for CodeRunnerTool { fn metadata(self) - ToolMetadata { ToolMetadata { name: run_code.to_string(), description: 在安全沙箱中执行 Python 代码.to_string(), parameters: vec![ ToolParamSchema { name: code.to_string(), param_type: string.to_string(), description: 要执行的 Python 代码.to_string(), required: true, }, ], } } async fn execute(self, params: serde_json::Value) - ToolResult { let code match params.get(code).and_then(|v| v.as_str()) { Some(c) c.to_string(), None return ToolResult { success: false, content: String::new(), data: None, error: Some(缺少必填参数: code.to_string()), }, }; // 安全沙箱检查禁止危险模块 let forbidden [os, subprocess, sys, shutil, importlib]; for module in forbidden { if code.contains(format!(import {}, module)) || code.contains(format!(from {}, module)) { return ToolResult { success: false, content: String::new(), data: None, error: Some(format!(禁止使用模块: {}, module)), }; } } // 模拟执行实际应使用子进程 沙箱 ToolResult { success: true, content: format!(代码执行结果: [输出占位符]\n代码长度: {} 字符, code.len()), data: Some(serde_json::json!({ code_len: code.len(), timeout: self.timeout })), error: None, } } }生产环境实战经验实现Tooltrait 时一个容易忽视的问题是execute的超时控制。网络搜索工具可能因为第三方 API 卡住代码执行工具可能写死循环。trait 本身没有超时语义需要在调用层加保护。我的做法是给ToolRegistry::call_tool套一层tokio::time::timeout超时就返回专门的错误不拖死整个 Agent。执行时间上限也写入工具元信息AI 看到后可以决定是否调用。这样 trait 的设计不需要变化但安全性明显提升。四、实现可插拔的工具注册表有了统一的Tooltrait工具注册表就很简单了——本质上就是一个HashMapString, Boxdyn Tooluse std::sync::Arc; use tokio::sync::RwLock; // 工具注册表 /// 工具注册表 —— 管理所有可用的工具 /// /// 使用 ArcRwLock 支持多 Agent 共享 运行时增删工具。 pub struct ToolRegistry { /// 工具名 → 工具实例 tools: RwLockHashMapString, Arcdyn Tool, } impl ToolRegistry { /// 创建空的工具注册表 pub fn new() - Self { ToolRegistry { tools: RwLock::new(HashMap::new()), } } /// 注册一个工具 pub async fn register(self, tool: impl Tool static) { let metadata tool.metadata(); let name metadata.name.clone(); let mut tools self.tools.write().await; tools.insert(name.clone(), Arc::new(tool)); println!(✅ 工具已注册: {}, name); } /// 注销一个工具 pub async fn unregister(self, name: str) - bool { let mut tools self.tools.write().await; let removed tools.remove(name).is_some(); if removed { println!(️ 工具已注销: {}, name); } removed } /// 获取所有工具的元信息列表 —— 用于构建 AI 的 system prompt pub async fn get_all_metadata(self) - VecToolMetadata { let tools self.tools.read().await; tools.values().map(|t| t.metadata()).collect() } /// 调用指定工具 pub async fn call_tool( self, name: str, params: serde_json::Value, ) - ResultToolResult, String { let tools self.tools.read().await; let tool tools .get(name) .ok_or_else(|| format!(工具不存在: {name}))?; // 先验证参数 tool.validate_params(params) .map_err(|e| format!(参数验证失败: {e}))?; // 执行工具 let result tool.execute(params).await; Ok(result) } /// 获取已注册工具的数量 pub async fn count(self) - usize { self.tools.read().await.len() } /// 获取所有工具名称列表 pub async fn tool_names(self) - VecString { let tools self.tools.read().await; tools.keys().cloned().collect() } } impl Default for ToolRegistry { fn default() - Self { Self::new() } } // 构建 AI Prompt 的辅助函数 /// 将工具元信息列表转换为 AI 可理解的 JSON 格式 pub fn tools_to_prompt_json(tools: [ToolMetadata]) - String { serde_json::to_string_pretty(tools).unwrap_or_else(|_| [].to_string()) } /// 构建带有工具列表的 system prompt pub fn build_system_prompt(tools: [ToolMetadata]) - String { let tools_json tools_to_prompt_json(tools); format!( r#你是一个有用的 AI 助手可以调用以下工具 {tools_json} 当需要使用工具时请返回 JSON 格式 {tool_call: {{name: 工具名, params: {{参数名: 参数值}}}} 如果不使用工具直接回复文本。#, ) }并发调用时的竞态问题RwLock保护注册表只在单 Agent 场景够用。多 Agent 并发调用同一个工具时要考虑工具是否有内部状态。比如文件读取工具同时操作一个临时文件可能读到一半被改掉。我的处理方式是让工具实现Clone每次调用从Arc克隆一个独立实例。对于需要全局状态的工具比如缓存单独用一个Mutex保护状态和注册表的锁分开。两个锁混在一起容易出现死锁画一张锁依赖图能避免很多问题。五、总结这篇文章用 Rust trait 实现了一个可插拔的 AI Agent 工具注册系统Tooltrait 定义统一接口metadata()告诉 AI 这个工具是干什么的execute()实际执行。所有工具实现同一个 traitAgent 可以用统一方式调用。ToolRegistry管理工具的生命周期注册、注销、查找、调用四个方法覆盖所有操作。用RwLock保护支持并发读写。类型安全 可插拔编译期保证工具接口一致运行时可以随时注册新工具。新增工具只需要写一个impl Tool for MyNewTool的文件不改核心代码。安全边界在工具内部ReadFileTool的路径沙箱、CodeRunnerTool的模块黑名单安全逻辑和注册表解耦每个工具负责自己的安全。作为独立开发者实现这个系统的最大收获是理解了面向接口编程的真正价值。不是教科书上说的方便扩展而是当你把接口定义清楚的瞬间剩下的代码就像填空一样自然。希望这篇文章对你在写 AI Agent 时有所帮助。有问题欢迎评论区交流