用 Rust 从零构建 AI Agent（三）：用 Trait 重构 Skill Registry

match name {
    "read_file" => read_file(args).await,
    "list_files" => list_files(args).await,
    _ => Err(...),
}

#[async_trait]
trait Skill: Send + Sync + 'static {
    type Input: DeserializeOwned + JsonSchema + Send + 'static;

    fn name(&self) -> &str;
    fn description(&self) -> &str;
    fn is_read_only(&self) -> bool { true }
    fn is_concurrency_safe(&self) -> bool { self.is_read_only() }

    async fn run(&self, input: Self::Input) -> Result<String, SkillError>;
}

#[derive(Deserialize, JsonSchema)]
struct ReadFileArgs {
    /// 文件路径，相对于项目根目录。
    path: String,
}

struct ReadFile;

#[async_trait]
impl Skill for ReadFile {
    type Input = ReadFileArgs;
    fn name(&self) -> &str { "read_file" }
    fn description(&self) -> &str {
        "Read a UTF-8 text file from the current project."
    }
    async fn run(&self, input: ReadFileArgs) -> Result<String, SkillError> {
        // 实际读取 + 沙箱校验
    }
}

let schema = schemars::schema_for!(S::Input);

{
  "type": "object",
  "properties": {
    "path": {
      "type": "string",
      "description": "文件路径，相对于项目根目录。"
    }
  },
  "required": ["path"]
}

#[async_trait]
trait DynSkill: Send + Sync {
    fn definition(&self) -> ToolDefinition;
    fn is_concurrency_safe(&self) -> bool;
    async fn run(&self, args: Value) -> Result<String, SkillError>;
}

#[async_trait]
impl<S: Skill> DynSkill for S {
    fn definition(&self) -> ToolDefinition {
        let schema = schema_for!(S::Input);
        ToolDefinition {
            name: self.name().to_string(),
            description: self.description().to_string(),
            input_schema: schema,
        }
    }

    fn is_concurrency_safe(&self) -> bool {
        self.is_concurrency_safe()
    }

    async fn run(&self, args: Value) -> Result<String, SkillError> {
        let input: S::Input = serde_json::from_value(args)?;
        self.run(input).await
    }
}

async fn run_many(
    &self,
    calls: &[(String, String, Value)], // (id, name, args)
) -> Vec<(String, Result<String, SkillError>)> {
    let mut parallel = Vec::new();
    let mut serial = Vec::new();

    for (id, name, args) in calls {
        match self.index.get(name) {
            Some(&i) if self.skills[i].is_concurrency_safe() => {
                parallel.push(self.run_one(id.clone(), i, args.clone()));
            }
            Some(&i) => {
                serial.push(self.run_one(id.clone(), i, args.clone()));
            }
            None => { /* 未知工具 */ }
        }
    }

    let mut results = join_all(parallel).await;
    for fut in serial {
        results.push(fut.await);
    }
    results
}

#[derive(Debug, Error)]
enum SkillError {
    #[error("invalid arguments: {0}")]
    BadArgs(#[from] serde_json::Error),
    #[error("safety check failed: {0}")]
    Safety(String),
    #[error("transient failure (worth retrying): {0}")]
    Transient(String),
    #[error("tool failed: {0}")]
    Tool(String),
}

async fn with_retry<F, Fut, T>(policy: RetryPolicy, mut op: F) -> Result<T, SkillError>
where
    F: FnMut() -> Fut,
    Fut: std::future::Future<Output = Result<T, SkillError>>,
{
    let mut delay = policy.initial_delay;
    for attempt in 0..policy.max_attempts {
        match op().await {
            Ok(v) => return Ok(v),
            Err(e) if e.is_transient() && attempt + 1 < policy.max_attempts => {
                tokio::time::sleep(delay).await;
                delay = delay.mul_f64(policy.backoff);
            }
            Err(e) => return Err(e),
        }
    }
    unreachable!()
}

#[async_trait]
pub trait SkillHook: Send + Sync + 'static {
    async fn before(
        &self,
        skill_name: &str,
        metadata: SkillMetadata,
        args: &Value,
    ) -> Result<HookOutcome, SkillError> { Ok(HookOutcome::Proceed) }

    async fn after(
        &self,
        skill_name: &str,
        args: &Value,
        result: Result<String, SkillError>,
    ) -> Result<String, SkillError> { result }
}

enum HookOutcome {
    Proceed,
    Modify(Value),
    Deny(String),
    Replace(String),
}

let mut registry = Registry::new();
registry
    .add(ListFiles)
    .add(ReadFile)
    .add(DeleteFile)
    .with_hook(PermissionGate::new(|_, md: SkillMetadata| !md.is_destructive));

fn search_hint(&self) -> Option<&str> { None }
fn should_defer(&self) -> bool { false }
fn aliases(&self) -> &[&'static str] { &[] }

impl Skill for DeleteFile {
    type Input = DeleteFileArgs;
    fn name(&self) -> &str { "delete_file" }
    fn description(&self) -> &str { "Delete a file from the project." }
    fn should_defer(&self) -> bool { true }
    fn search_hint(&self) -> Option<&str> { Some("remove file rm" ) }
    fn aliases(&self) -> &[&'static str] { &["remove_file"] }
    // ...
}

let mut registry = Registry::new();
registry.add(ListFiles).add(ReadFile);

let tools: Vec<Tool> = registry
    .definitions()
    .into_iter()
    .map(Into::into)
    .collect();

// ... 构建系统提示词 ...

for _ in 0..MAX_TURNS {
    let response = with_retry(policy, || async {
        send(&http, &api_key, &system_blocks, &tools, &messages).await
    }).await?;

    // 把 assistant 消息追加到 messages
    // 如果包含 tool_use 块，调用 registry.run_many，得到 tool_result
    // 如果包含最终回复，返回给用户
}