Tool trait 的双 trait 设计

Tool trait 的双 trait 设计

本节摘要:工具系统的根基是一个叫 Tool 的 trait——它定义了「一个工具长什么样」。但 Rust 的 trait 有一个棘手的限制:要塞进同一个集合(动态分发),trait 必须是「对象安全」的,而带关联类型、泛型、impl Trait 返回值的 trait 不是对象安全的。Grok Build 的解法是一个精巧的「双 trait + 毯子 impl」设计:让工具作者写一个用户友好的泛型 Tool trait,框架内部维护一个对象安全的 ToolDyn 镜像,再用毯子 impl 把它们连起来。本节会拆解这个设计,讲清它解决了什么问题、为什么这样设计,以及它带来的好处。这是本章最「Rust 味」的一节,但理解它,你就掌握了 Grok Build 工具系统的灵魂。

一、工具系统的接口需求

先看工具系统需要什么样的接口。一个工具,从最朴素的角度,需要能:

• 报告自己的身份(名字、ID)
• 描述自己(给模型看的说明、参数 JSON Schema)
• 执行:接收参数,返回结果

用 Rust 的 trait 描述,大致是:

trait Tool {
    type Args;            # 参数类型
    type Output;          # 返回值类型
    
    fn id(&self) -> ToolId;
    fn description(&self) -> ToolDescription;
    fn execute(&self, args: Self::Args) -> Self::Output;
}
​

这是「理想形态」——每个工具有自己强类型的 Args 与 Output,写起来直观、有类型检查。

二、对象安全的难题

理想形态很好,但遇到一个现实难题:不同的工具有不同的 Args 与 Output 类型,怎么把它们塞进同一个集合?

假设工具注册表要这样:

let tools: Vec<Box<dyn Tool>> = vec![
    Box::new(ReadFileTool),
    Box::new(BashTool),
    Box::new(GrepTool),
]
​

问题来了:dyn Tool 要求 Tool 是「对象安全」的(object-safe)。Rust 的对象安全规则大致是:trait 方法不能返回 Self、不能有泛型方法、不能有关联类型(除非用 where Self::Args: ... 约束)。

上面的 Tool trait 有 type Args 和 type Output,不是对象安全的,不能做 dyn Tool。

为什么对象安全要求这样

对象安全的存在,是因为 dyn Trait 是「类型擦除」——编译器要把不同具体类型(ReadFileTool、BashTool 等)统一成一个胖指针(数据指针 + 虚表)。虚表里要能放下每个方法的函数指针。但如果方法签名依赖具体类型(如返回 Self::Output),编译器无法为不同 Output 类型生成统一的虚表入口。所以要求方法签名「不依赖具体类型」。

三、朴素方案的局限

有几种朴素的应对,但都不理想:

方案 A:去掉关联类型,全用 serde_json::Value

trait Tool {
    fn execute(&self, args: serde_json::Value) -> serde_json::Value;
}
​

这样可以对象安全。但代价是失去类型安全——工具作者要在 execute 内部手动 serde_json::from_value(args) 解析参数、手动构造返回值的 JSON。每个工具都写一遍这种样板代码,既啰嗦又容易出错。

方案 B:每个工具自己存,不进统一集合

不做 Vec<Box<dyn Tool>>,而是用枚举或 match 分发。但工具数量多(50+),且要支持外部注册(MCP、插件),枚举方案扩展性差,新增工具要改枚举与 match,不符合开闭原则。

这两种都不理想。Grok Build 的解法是第三种:双 trait + 毯子 impl。

四、双 trait 设计

Grok Build 的方案(在 xai-tool-runtime/src/tool.rs):

第一个 trait:Tool(用户友好,非对象安全)

trait Tool: Send + Sync {
    type Args: Deserialize + JsonSchema;       # 参数类型,带反序列化与 Schema 生成
    type Output: Serialize + ToolOutput;        # 返回值类型,带序列化
    
    fn id(&self) -> ToolId;
    fn description(&self, ctx) -> ToolDescription;
    
    # 执行:返回一条流(ToolStream,见下一节)
    fn execute(&self, ctx, args: Self::Args) -> ToolStream<Self::Output>;
    
    # 或者简单的同步版本(默认实现把它包成单条流)
    fn run(&self, ctx, args: Self::Args) -> Result<Self::Output, ToolError> {
        NotImplemented
    }
}
​

特点:

• 带 type Args 和 type Output——工具作者写强类型代码
• Args: Deserialize + JsonSchema——参数能从 JSON 反序列化,且能自动生成 JSON Schema 给模型
• execute 返回 ToolStream<Self::Output>——流式执行(下一节详谈)
• 这个 trait 不是对象安全(有关联类型),不能 dyn Tool

工具作者只用关心这个 trait——写自己的 Args/Output 类型,实现 execute。享受类型安全,无需关心动态分发。

第二个 trait:ToolDyn(对象安全,框架内部)

trait ToolDyn: Send + Sync {
    fn id(&self) -> ToolId;
    fn description_dyn(&self, ctx) -> ToolDescription;
    
    async fn execute_dyn(
        &self,
        ctx,
        args: serde_json::Value      # 类型擦除的参数
    ) -> ToolStream<TypedToolOutput>;  # 类型擦除的返回
}

type ArcTool = Arc<dyn ToolDyn>;
​

特点:

• 没有关联类型——对象安全,可以 dyn ToolDyn
• 方法签名用 serde_json::Value 与 TypedToolOutput(类型擦除的统一返回)替代强类型
• 注册表实际存的是 Arc<dyn ToolDyn>

这个 trait 是框架内部的,工具作者通常不直接实现。

毯子 impl:连接两个 trait

impl<T: Tool> ToolDyn for T {
    fn id(&self) -> ToolId {
        Tool::id(self)                    # 直接转发
    }
    
    fn description_dyn(&self, ctx) -> ToolDescription {
        Tool::description(self, ctx)      # 直接转发
    }
    
    async fn execute_dyn(&self, ctx, args: serde_json::Value) -> ToolStream<TypedToolOutput> {
        # ① 把类型擦除的 Value 反序列化成工具的强类型 Args
        let typed_args: T::Args = serde_json::from_value(args)?;
        
        # ② 调用强类型 execute
        let stream = Tool::execute(self, ctx, typed_args);
        
        # ③ 把强类型 Output 流,映射成类型擦除的 TypedToolOutput 流
        stream.map(|item| match item {
            Progress(p) => Progress(p.into_typed()),
            Terminal(result) => Terminal(result.into_typed()),
        })
    }
}
​

毯子 impl 的作用:为所有 T: Tool 自动实现 ToolDyn。这样:

• 工具作者实现 Tool(强类型)
• 框架自动获得它的 ToolDyn 实现(类型擦除)
• 注册表存 Arc<dyn ToolDyn>,动态分发

关键概念:双 trait 设计的核心是「关注点分离」——工具作者关心类型安全与易用性(用 Tool),框架关心动态分发与统一存储(用 ToolDyn),毯子 impl 把两者桥接。作者写一次强类型代码,自动获得动态分发能力,无需手写任何类型擦除的样板。

五、为什么这是好设计

把这个设计放在一起评估,它好在几个方面:

好处一:工具作者体验好

工具作者只实现 Tool trait,享受完整的类型安全:

• Args 与 Output 是强类型,编译期检查
• 参数从 JSON 反序列化自动完成(不用手写 from_value)
• 返回值序列化自动完成
• JSON Schema 自动生成(给模型看)

对比「方案 A 全用 Value」,这省去了大量样板代码与潜在错误。

好处二:框架获得动态分发

注册表可以统一存 Arc<dyn ToolDyn>,支持任意数量、任意类型的工具,包括运行时注册的(MCP、插件)。这符合开闭原则——新增工具不必改框架代码。

好处三:桥接代码只写一次

毯子 impl 把强类型与类型擦除的转换集中在一处。转换逻辑(反序列化、序列化、Schema 生成)只写一次,所有工具复用。这避免了每个工具都重复写转换代码。

好处四:扩展性好

新增一个工具,只需:

struct MyTool;

impl Tool for MyTool {
    type Args = MyArgs;
    type Output = MyOutput;
    
    fn id(&self) -> ToolId { ... }
    fn description(&self, ctx) -> ToolDescription { ... }
    fn execute(&self, ctx, args: Self::Args) -> ToolStream<Self::Output> { ... }
}
​

注册时:

registry.register::<MyTool>()
​

框架自动获得它的 ToolDyn,塞进 ArcTool 集合。无需任何额外代码。

六、Args 的 JsonSchema 约束

注意 Tool trait 里 type Args: Deserialize + JsonSchema 的 JsonSchema 约束。这不是装饰——它让框架能为每个工具的参数自动生成 JSON Schema,发给模型。

为什么需要 Schema

模型不能直接调用 Rust 函数,它生成的是 JSON。框架要告诉模型「这个工具接受什么参数」,用 JSON Schema 描述:

ReadFile 工具的 Schema(简化):
{
    "type": "object",
    "properties": {
        "path": { "type": "string", "description": "要读取的文件路径" }
    },
    "required": ["path"]
}
​

模型看到这个 Schema,就知道「调用 ReadFile 要传一个 path 字符串」。它生成形如 {"path": "src/main.rs"} 的 JSON,框架反序列化成 ReadFile 的 Args。

自动生成

通过 #[derive(JsonSchema)](schemars 库),Rust 结构体自动生成 Schema。工具作者定义 Args 结构体,加 derive,框架自动拿到 Schema——无需手写。

这是「类型即文档」的体现:你定义的 Args 类型,既是编译期的类型检查,也是运行期发给模型的说明。改了类型,Schema 自动更新,不会出现「文档与代码不一致」。

七、ToolOutput 与结果标准化

类似地,type Output: Serialize + ToolOutput 里的 ToolOutput 约束,让返回值能被「标准化」成统一的呈现。

工具的 Output 是强类型(如 ReadFile 可能返回 FileContent { path, content, lines }),但塞回历史给模型看时,需要转成统一的文本形式。ToolOutput trait 提供这个转换:

trait ToolOutput {
    fn to_tool_result(&self) -> ToolResult {
        # 默认实现:把 self 序列化成 JSON 字符串或格式化文本
    }
}
​

这样无论工具的 Output 是什么类型,最终都变成一个标准的 ToolResult(通常是文本),塞回会话历史。模型看到的是统一的「工具结果消息」,不必关心工具内部的 Output 类型。

八、设计的整体意义

把双 trait 设计放回更大的语境,它体现了几个值得借鉴的工程原则:

原则一:关注点分离

用户的关心点(写强类型、易用的工具)与框架的关心点(动态分发、统一存储)是正交的。双 trait 让它们各自有专门的服务接口。

原则二:类型擦除集中化

类型擦除(serde_json::Value、TypedToolOutput)是「不可避免但应集中」的复杂度。毯子 impl 把它集中在一处,而不是散落到每个工具。

原则三:自动化优于手写

Schema 生成、参数反序列化、结果标准化,都通过 derive 与默认实现自动化。工具作者只写业务逻辑,样板交给框架。

原则四:开闭原则

对扩展开放(新增工具不改框架)、对修改关闭(框架核心稳定)。这是大型工具系统长期演进的关键。

设计警示:双 trait 设计不是「为了炫技」,而是 Rust 类型系统的现实约束下的务实解法。在支持高质量动态分发(如 Go 的接口、Java 的反射)的语言里,这种复杂度可以更低;但在 Rust 里,这种设计换来的是零成本抽象与编译期安全,代价是设计本身的精巧。理解它,你就理解了 Grok Build 工具系统的工程深度。

本节要点回顾

1. 理想工具接口带关联类型:Args/Output 强类型,但非对象安全,不能 dyn Tool。
2. 对象安全要求:trait 方法不能依赖具体类型(关联类型、Self、泛型方法),否则虚表无法统一。
3. 朴素方案的局限:全用 Value 失去类型安全;枚举分发扩展性差。
4. 双 trait 设计:Tool(用户友好,强类型,非对象安全)+ ToolDyn(对象安全,类型擦除,框架内部)。
5. 毯子 impl 桥接:为所有 T: Tool 自动实现 ToolDyn,类型擦除集中在一处。
6. 好处:作者体验好(强类型)、框架获动态分发、桥接代码只写一次、扩展性好。
7. Args 的 JsonSchema 约束:自动生成参数 Schema 发给模型,「类型即文档」。
8. ToolOutput 标准化:强类型返回值转成统一 ToolResult 文本塞回历史。
9. 体现的工程原则:关注点分离、类型擦除集中化、自动化优于手写、开闭原则。

下一节,我们看 ToolStream——工具执行不是返回单个值,而是返回一条「Progress* then Terminal」的流,这是流式执行的统一形态。