agent_diva_agent/agent_loop/
loop_tools.rs

1use super::{AgentLoop, ToolConfig};
2use crate::tool_config::network::NetworkToolConfig;
3use agent_diva_core::config::MCPServerConfig;
4use agent_diva_core::security::{SecurityConfig, SecurityLevel, SecurityPolicy};
5use agent_diva_tools::{
6    load_mcp_tools_sync, CronTool, EditFileTool, ExecTool, ListDirTool, ReadAttachmentTool,
7    ReadFileTool, SpawnTool, ToolError, ToolRegistry, WebFetchTool, WebSearchTool, WriteFileTool,
8};
9use std::collections::HashMap;
10use std::sync::Arc;
11use tracing::info;
12
13impl AgentLoop {
14    /// Register default tools (for use after construction)
15    pub fn register_default_tools(&mut self, tool_config: ToolConfig) {
16        // Register spawn tool
17        let sm = self.subagent_manager.clone();
18        self.tools.register(Arc::new(SpawnTool::new(
19            move |task, label, channel, chat_id| {
20                let sm = sm.clone();
21                async move {
22                    sm.spawn(task, label, channel, chat_id)
23                        .await
24                        .map_err(|e| ToolError::ExecutionFailed(e.to_string()))
25                }
26            },
27        )));
28
29        // Register file system tools with SecurityPolicy
30        let security_config = if tool_config.restrict_to_workspace {
31            SecurityConfig {
32                level: SecurityLevel::Standard,
33                workspace_only: true,
34                ..SecurityConfig::default()
35            }
36        } else {
37            SecurityConfig::default()
38        };
39        let security = Arc::new(SecurityPolicy::with_config(
40            self.workspace.clone(),
41            security_config,
42        ));
43        self.tools
44            .register(Arc::new(ReadFileTool::new(security.clone())));
45        self.tools
46            .register(Arc::new(WriteFileTool::new(security.clone())));
47        self.tools
48            .register(Arc::new(EditFileTool::new(security.clone())));
49        self.tools.register(Arc::new(ListDirTool::new(security)));
50
51        // Register attachment read tool with shared FileManager
52        self.tools
53            .register(Arc::new(ReadAttachmentTool::new(self.file_manager.clone())));
54
55        // Register shell tool
56        self.tools.register(Arc::new(ExecTool::with_config(
57            tool_config.exec_timeout,
58            Some(self.workspace.clone()),
59            tool_config.restrict_to_workspace,
60        )));
61
62        // Register web tools
63        Self::register_web_tools(&mut self.tools, &tool_config.network);
64
65        // Register MCP tools discovered from configured servers
66        for mcp_tool in load_mcp_tools_sync(&tool_config.mcp_servers) {
67            self.tools.register(mcp_tool);
68        }
69
70        // Register cron tool when scheduling is configured
71        if let Some(cron_service) = tool_config.cron_service {
72            self.tools.register(Arc::new(CronTool::new(cron_service)));
73        }
74    }
75
76    pub(super) fn register_web_tools(tools: &mut ToolRegistry, network: &NetworkToolConfig) {
77        if network.web.search.enabled {
78            tools.register(Arc::new(WebSearchTool::with_provider_and_max_results(
79                network.web.search.provider.clone(),
80                network.web.search.api_key.clone(),
81                network.web.search.normalized_max_results(),
82            )));
83        }
84        if network.web.fetch.enabled {
85            tools.register(Arc::new(WebFetchTool::new()));
86        }
87    }
88
89    pub(super) async fn apply_network_config(&mut self, network: NetworkToolConfig) {
90        self.tools.unregister("web_search");
91        self.tools.unregister("web_fetch");
92        Self::register_web_tools(&mut self.tools, &network);
93        self.subagent_manager.update_network_config(network).await;
94        info!("Applied runtime network tool configuration update");
95    }
96
97    pub(super) async fn apply_mcp_config(&mut self, servers: HashMap<String, MCPServerConfig>) {
98        for name in self
99            .tools
100            .tool_names()
101            .into_iter()
102            .filter(|name| name.starts_with("mcp_"))
103        {
104            self.tools.unregister(&name);
105        }
106
107        for mcp_tool in load_mcp_tools_sync(&servers) {
108            self.tools.register(mcp_tool);
109        }
110
111        info!("Applied runtime MCP tool configuration update");
112    }
113}
agent_diva_agent/agent_loop/loop_tools.rs

agent_diva_agent/agent_loop/
loop_tools.rs
