agent_diva_agent/
tool_assembly.rs

1use crate::tool_config::{builtin::BuiltInToolsConfig, network::NetworkToolConfig};
2use agent_diva_core::config::MCPServerConfig;
3use agent_diva_core::cron::CronService;
4use agent_diva_core::security::{SecurityConfig, SecurityLevel, SecurityPolicy};
5use agent_diva_files::FileManager;
6use agent_diva_tooling::{Tool, ToolError, ToolRegistry};
7use agent_diva_tools::{
8    load_mcp_tools_sync, CronTool, EditFileTool, ExecTool, ListDirTool, ReadAttachmentTool,
9    ReadFileTool, SpawnTool, WebFetchTool, WebSearchTool, WriteFileTool,
10};
11use std::collections::HashMap;
12use std::path::PathBuf;
13use std::sync::Arc;
14
15#[async_trait::async_trait]
16pub trait SubagentSpawner: Send + Sync {
17    async fn spawn(
18        &self,
19        task: String,
20        label: Option<String>,
21        channel: String,
22        chat_id: String,
23    ) -> Result<String, ToolError>;
24}
25
26pub struct ToolAssembly {
27    workspace: PathBuf,
28    builtin_config: BuiltInToolsConfig,
29    network_config: NetworkToolConfig,
30    exec_timeout: u64,
31    restrict_to_workspace: bool,
32    mcp_servers: HashMap<String, MCPServerConfig>,
33    cron_service: Option<Arc<CronService>>,
34    custom_tools: Vec<Arc<dyn Tool>>,
35    subagent_spawner: Option<Arc<dyn SubagentSpawner>>,
36    file_manager: Option<Arc<FileManager>>,
37}
38
39impl ToolAssembly {
40    pub fn new(workspace: PathBuf) -> Self {
41        Self {
42            workspace,
43            builtin_config: BuiltInToolsConfig::default(),
44            network_config: NetworkToolConfig::default(),
45            exec_timeout: 60,
46            restrict_to_workspace: false,
47            mcp_servers: HashMap::new(),
48            cron_service: None,
49            custom_tools: Vec::new(),
50            subagent_spawner: None,
51            file_manager: None,
52        }
53    }
54
55    pub fn builtin(mut self, config: BuiltInToolsConfig) -> Self {
56        self.builtin_config = config;
57        self
58    }
59
60    pub fn with_network_config(mut self, config: NetworkToolConfig) -> Self {
61        self.network_config = config;
62        self
63    }
64
65    pub fn with_exec_timeout(mut self, timeout: u64) -> Self {
66        self.exec_timeout = timeout;
67        self
68    }
69
70    pub fn restrict_to_workspace(mut self, restrict: bool) -> Self {
71        self.restrict_to_workspace = restrict;
72        self
73    }
74
75    pub fn mcp_servers(mut self, servers: HashMap<String, MCPServerConfig>) -> Self {
76        self.mcp_servers = servers;
77        self
78    }
79
80    pub fn with_cron_service(mut self, service: Arc<CronService>) -> Self {
81        self.cron_service = Some(service);
82        self
83    }
84
85    pub fn with_subagent_spawner(mut self, spawner: Arc<dyn SubagentSpawner>) -> Self {
86        self.subagent_spawner = Some(spawner);
87        self
88    }
89
90    pub fn with_file_manager(mut self, file_manager: Arc<FileManager>) -> Self {
91        self.file_manager = Some(file_manager);
92        self
93    }
94
95    pub fn with_tool(mut self, tool: Arc<dyn Tool>) -> Self {
96        self.custom_tools.push(tool);
97        self
98    }
99
100    pub fn with_tools(mut self, tools: Vec<Arc<dyn Tool>>) -> Self {
101        self.custom_tools.extend(tools);
102        self
103    }
104
105    pub fn build(self) -> ToolRegistry {
106        self.build_internal(false)
107    }
108
109    pub fn build_subagent_registry(mut self) -> ToolRegistry {
110        self.builtin_config = self.builtin_config.for_subagent();
111        self.subagent_spawner = None;
112        self.cron_service = None;
113        self.file_manager = None;
114        self.build_internal(true)
115    }
116
117    fn build_internal(self, subagent_mode: bool) -> ToolRegistry {
118        let mut registry = ToolRegistry::new();
119
120        if self.builtin_config.filesystem {
121            let security_config = if self.restrict_to_workspace {
122                SecurityConfig {
123                    level: SecurityLevel::Standard,
124                    workspace_only: true,
125                    ..SecurityConfig::default()
126                }
127            } else {
128                SecurityConfig::default()
129            };
130            let security = Arc::new(SecurityPolicy::with_config(
131                self.workspace.clone(),
132                security_config,
133            ));
134            registry.register(Arc::new(ReadFileTool::new(security.clone())));
135            registry.register(Arc::new(WriteFileTool::new(security.clone())));
136            registry.register(Arc::new(EditFileTool::new(security.clone())));
137            registry.register(Arc::new(ListDirTool::new(security)));
138        }
139
140        if self.builtin_config.attachment {
141            if let Some(file_manager) = self.file_manager {
142                registry.register(Arc::new(ReadAttachmentTool::new(file_manager)));
143            }
144        }
145
146        if self.builtin_config.shell {
147            registry.register(Arc::new(ExecTool::with_config(
148                self.exec_timeout,
149                Some(self.workspace.clone()),
150                self.restrict_to_workspace,
151            )));
152        }
153
154        if self.builtin_config.web_search && self.network_config.web.search.enabled {
155            registry.register(Arc::new(WebSearchTool::with_provider_and_max_results(
156                self.network_config.web.search.provider.clone(),
157                self.network_config.web.search.api_key.clone(),
158                self.network_config.web.search.normalized_max_results(),
159            )));
160        }
161
162        if self.builtin_config.web_fetch && self.network_config.web.fetch.enabled {
163            registry.register(Arc::new(WebFetchTool::new()));
164        }
165
166        if self.builtin_config.spawn && !subagent_mode {
167            if let Some(spawner) = self.subagent_spawner {
168                registry.register(Arc::new(SpawnTool::new(
169                    move |task, label, channel, chat_id| {
170                        let spawner = spawner.clone();
171                        async move { spawner.spawn(task, label, channel, chat_id).await }
172                    },
173                )));
174            }
175        }
176
177        if self.builtin_config.mcp && !self.mcp_servers.is_empty() {
178            for tool in load_mcp_tools_sync(&self.mcp_servers) {
179                registry.register(tool);
180            }
181        }
182
183        if self.builtin_config.cron && !subagent_mode {
184            if let Some(cron_service) = self.cron_service {
185                registry.register(Arc::new(CronTool::new(cron_service)));
186            }
187        }
188
189        for tool in self.custom_tools {
190            registry.register(tool);
191        }
192
193        registry
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200
201    #[test]
202    fn test_tool_assembly_minimal() {
203        let registry = ToolAssembly::new(PathBuf::from("/tmp/test"))
204            .builtin(BuiltInToolsConfig::minimal())
205            .build();
206
207        assert!(registry.has("read_file"));
208        assert!(registry.has("write_file"));
209        assert!(registry.has("edit_file"));
210        assert!(registry.has("list_dir"));
211        assert!(!registry.has("exec"));
212        assert!(!registry.has("web_search"));
213        assert!(!registry.has("web_fetch"));
214    }
215
216    #[test]
217    fn test_tool_assembly_none() {
218        let registry = ToolAssembly::new(PathBuf::from("/tmp/test"))
219            .builtin(BuiltInToolsConfig::none())
220            .build();
221
222        assert!(registry.is_empty());
223    }
224
225    #[test]
226    fn test_tool_assembly_respects_split_web_flags() {
227        let registry = ToolAssembly::new(PathBuf::from("/tmp/test"))
228            .builtin(BuiltInToolsConfig {
229                web_search: true,
230                web_fetch: false,
231                ..BuiltInToolsConfig::none()
232            })
233            .with_network_config(NetworkToolConfig::default())
234            .build();
235
236        assert!(registry.has("web_search"));
237        assert!(!registry.has("web_fetch"));
238    }
239
240    #[test]
241    fn test_tool_assembly_subagent_mode_disables_spawn_and_attachment() {
242        let registry = ToolAssembly::new(PathBuf::from("/tmp/test"))
243            .builtin(BuiltInToolsConfig {
244                filesystem: true,
245                spawn: true,
246                attachment: true,
247                ..BuiltInToolsConfig::none()
248            })
249            .build_subagent_registry();
250
251        assert!(registry.has("read_file"));
252        assert!(!registry.has("spawn"));
253        assert!(!registry.has("read_attachment"));
254    }
255}
agent_diva_agent/tool_assembly.rs

agent_diva_agent/
tool_assembly.rs
