codei_tools/
registry.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use codei_config::ResolvedConfig;
5use serde_json::Value;
6
7use crate::approval::{ApprovalHandler, ApprovalRequest};
8use crate::ToolError;
9
10#[derive(Debug, Clone)]
11pub struct ToolResult {
12    pub content: String,
13    pub is_error: bool,
14}
15
16#[derive(Clone)]
17pub struct ToolContext {
18    pub cwd: std::path::PathBuf,
19    pub config: Arc<ResolvedConfig>,
20    pub approval: Arc<dyn ApprovalHandler>,
21}
22
23#[async_trait]
24pub trait Tool: Send + Sync {
25    fn name(&self) -> &str;
26    fn description(&self) -> &str;
27    fn parameters_schema(&self) -> Value;
28    fn requires_approval(&self) -> bool {
29        false
30    }
31
32    async fn execute(&self, ctx: &ToolContext, args: Value) -> Result<ToolResult, ToolError>;
33}
34
35pub struct ToolRegistry {
36    tools: Vec<Box<dyn Tool>>,
37}
38
39impl ToolRegistry {
40    pub fn new() -> Self {
41        Self { tools: Vec::new() }
42    }
43
44    pub fn register(&mut self, tool: Box<dyn Tool>) {
45        self.tools.push(tool);
46    }
47
48    pub fn definitions(&self) -> Vec<codei_llm::ToolDefinition> {
49        self.tools
50            .iter()
51            .map(|tool| codei_llm::ToolDefinition {
52                name: tool.name().to_string(),
53                description: tool.description().to_string(),
54                parameters: tool.parameters_schema(),
55            })
56            .collect()
57    }
58
59    pub fn get(&self, name: &str) -> Option<&dyn Tool> {
60        self.tools
61            .iter()
62            .find(|t| t.name() == name)
63            .map(|t| t.as_ref())
64    }
65
66    pub async fn execute(
67        &self,
68        ctx: &ToolContext,
69        name: &str,
70        args: Value,
71    ) -> Result<ToolResult, ToolError> {
72        let tool = self.get(name).ok_or_else(|| ToolError::Failed {
73            name: name.to_string(),
74            message: "unknown tool".into(),
75        })?;
76
77        if tool.requires_approval() {
78            let response = ctx
79                .approval
80                .approve(ApprovalRequest {
81                    tool_name: name.to_string(),
82                    arguments: args.clone(),
83                })
84                .await;
85            if !response.approved {
86                return Err(ToolError::Denied);
87            }
88        }
89
90        tool.execute(ctx, args).await
91    }
92}
93
94impl Default for ToolRegistry {
95    fn default() -> Self {
96        Self::new()
97    }
98}
codei_tools/registry.rs

codei_tools/
registry.rs
