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

agent_diva_agent/agent_loop/
loop_tools.rs
