tycode_core/tools/
registry.rs

1use crate::ai::{ToolDefinition, ToolUseData};
2use crate::tools::r#trait::{ToolCallHandle, ToolCategory, ToolExecutor, ToolRequest};
3use crate::tools::ToolName;
4use std::collections::BTreeMap;
5use std::sync::Arc;
6use tracing::{debug, error};
7
8pub struct ToolRegistry {
9    tools: BTreeMap<String, Arc<dyn ToolExecutor>>,
10}
11
12impl ToolRegistry {
13    pub fn new(tools: Vec<Arc<dyn ToolExecutor>>) -> Self {
14        let mut registry = Self {
15            tools: BTreeMap::new(),
16        };
17
18        for tool in tools {
19            registry.register_tool(tool);
20        }
21
22        registry
23    }
24
25    pub fn register_tool(&mut self, tool: Arc<dyn ToolExecutor>) {
26        let name = tool.name().to_string();
27        debug!(tool_name = %name, "Registering tool");
28        self.tools.insert(name, tool);
29    }
30
31    pub fn get_tool_definitions(&self, tool_names: &[ToolName]) -> Vec<ToolDefinition> {
32        let allowed_names: Vec<String> = tool_names.iter().map(|name| name.to_string()).collect();
33
34        self.tools
35            .iter()
36            .filter(|(name, _)| allowed_names.contains(name) || name.starts_with("mcp_"))
37            .map(|(_, tool)| ToolDefinition {
38                name: tool.name().to_string(),
39                description: tool.description().to_string(),
40                input_schema: tool.input_schema(),
41            })
42            .collect()
43    }
44
45    pub async fn process_tools(
46        &self,
47        tool_use: &ToolUseData,
48        allowed_tools: &[ToolName],
49    ) -> Result<Box<dyn ToolCallHandle>, String> {
50        let mut allowed_names: Vec<&str> = allowed_tools
51            .iter()
52            .map(|tool| tool.as_str())
53            .filter(|name| self.tools.contains_key(*name))
54            .collect();
55
56        // MCP tools are dynamically discovered and should be allowed for all agents
57        for name in self.tools.keys() {
58            if name.starts_with("mcp_") && !allowed_names.contains(&name.as_str()) {
59                allowed_names.push(name.as_str());
60            }
61        }
62
63        let tool = match self.tools.get(&tool_use.name) {
64            Some(tool) => tool,
65            None => {
66                let available = allowed_names.join(", ");
67                error!(tool_name = %tool_use.name, "Unknown tool");
68                return Err(format!(
69                    "Unknown tool: {}. Available tools: {}",
70                    tool_use.name, available
71                ));
72            }
73        };
74
75        if !allowed_names.contains(&tool_use.name.as_str()) {
76            debug!(
77                tool_name = %tool_use.name,
78                allowed_tools = ?allowed_names,
79                "Tool not in allowed list for current agent"
80            );
81            return Err(format!(
82                "Tool not available for current agent: {}",
83                tool_use.name
84            ));
85        }
86
87        let schema = tool.input_schema();
88        let coerced_arguments =
89            match crate::tools::fuzzy_json::coerce_to_schema(&tool_use.arguments, &schema) {
90                Ok(args) => args,
91                Err(e) => {
92                    error!(?e, tool_name = %tool_use.name, "Failed to coerce tool arguments");
93                    return Err(format!("Failed to coerce arguments: {e:?}"));
94                }
95            };
96
97        let request = ToolRequest::new(coerced_arguments, tool_use.id.clone());
98        tool.process(&request).await.map_err(|e| {
99            error!(?e, tool_name = %tool_use.name, "Tool processing failed");
100            format!("Error: {e:?}")
101        })
102    }
103
104    pub fn list_tools(&self) -> Vec<&str> {
105        self.tools.keys().map(|s| s.as_str()).collect()
106    }
107
108    /// Get tool executor by name
109    pub fn get_tool_executor_by_name(&self, name: &str) -> Option<&Arc<dyn ToolExecutor>> {
110        self.tools.get(name)
111    }
112
113    /// Get tool category by name
114    pub fn get_tool_category_by_name(&self, name: &str) -> Option<ToolCategory> {
115        self.tools.get(name).map(|executor| executor.category())
116    }
117}
tycode_core/tools/registry.rs

tycode_core/tools/
registry.rs
