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

agent_diva_agent/agent_loop/
loop_tools.rs
