返回 CodeWhale
tool_setup.rs
根目录 / crates / tui / src / core / engine / tool_setup.rs
1 //! Per-turn tool registry setup.
2 //!
3 //! This keeps mode/feature-specific registry construction out of the send path.
4
5 use super::*;
6 use crate::core::authority::shell_policy_for_mode;
7 use crate::tools::AgentToolSurfaceOptions;
8 use crate::worker_profile::ShellPolicy;
9
10 fn should_register_remember_tool(memory_enabled: bool) -> bool {
11 memory_enabled
12 }
13
14 impl Engine {
15 pub(super) fn agent_tool_surface_options(
16 &self,
17 shell_policy: ShellPolicy,
18 ) -> AgentToolSurfaceOptions {
19 let mut options = AgentToolSurfaceOptions::new(shell_policy);
20 options.apply_patch_enabled = self.config.features.enabled(Feature::ApplyPatch);
21 options.web_search_enabled = self.config.features.enabled(Feature::WebSearch);
22 options.memory_tool_enabled = should_register_remember_tool(self.config.memory_enabled);
23 options.vision_config = if self.config.features.enabled(Feature::VisionModel) {
24 self.config.vision_config.clone()
25 } else {
26 None
27 };
28 options.speech_output_dir = self.config.speech_output_dir.clone();
29 options.goal_state = Some(self.config.goal_state.clone());
30 options.verify_tool_enabled = self.config.features.enabled(Feature::Verify);
31 options
32 }
33
34 #[cfg(test)]
35 pub(super) fn build_turn_tool_registry_builder(
36 &self,
37 mode: AppMode,
38 todo_list: SharedTodoList,
39 plan_state: SharedPlanState,
40 ) -> ToolRegistryBuilder {
41 self.build_turn_tool_registry_builder_for_route(
42 mode,
43 self.session.allow_shell,
44 self.deepseek_client.clone(),
45 &self.session.model,
46 todo_list,
47 plan_state,
48 )
49 }
50
51 /// Build the registry from the route and authority already resolved for
52 /// this turn. Preview calls this before either is installed on the engine,
53 /// so reading `self.session` here would describe the previous turn's shell
54 /// posture, client, and model.
55 #[allow(clippy::too_many_arguments)]
56 pub(super) fn build_turn_tool_registry_builder_for_route(
57 &self,
58 mode: AppMode,
59 allow_shell: bool,
60 client: Option<DeepSeekClient>,
61 model: &str,
62 todo_list: SharedTodoList,
63 plan_state: SharedPlanState,
64 ) -> ToolRegistryBuilder {
65 let shell_policy = shell_policy_for_mode(mode, allow_shell);
66 if mode != AppMode::Plan {
67 let mut builder = ToolRegistryBuilder::new().with_agent_runtime_surface(
68 client.clone(),
69 model.to_string(),
70 self.agent_tool_surface_options(shell_policy),
71 todo_list,
72 plan_state,
73 );
74 if self.config.features.enabled(Feature::Mcp) {
75 builder = builder.with_registry_mcp_sync_tool();
76 }
77 // `start_mcp_server` belongs to every executable mode. Keep its
78 // handler aligned with the model catalog, which always loads the
79 // tool while MCP is enabled. The former early return registered
80 // it only in Plan mode, so Agent/Full Access advertised a tool
81 // that could never cross the execution boundary.
82 if let Some(ref pool) = self.mcp_pool {
83 builder = builder
84 .with_runtime_mcp_tool(Arc::clone(pool))
85 .with_registry_mcp_start_tool(Arc::clone(pool));
86 }
87 return builder;
88 }
89
90 let mut builder = {
91 let builder = ToolRegistryBuilder::new()
92 .with_read_only_file_tools()
93 .with_search_tools()
94 .with_git_tools()
95 .with_git_history_tools()
96 .with_diagnostics_tool()
97 .with_skill_tools()
98 .with_validation_tools()
99 .with_handle_tools()
100 .with_runtime_read_only_task_tools()
101 .with_todo_tool(todo_list)
102 .with_plan_tool(plan_state)
103 .with_goal_tools(self.config.goal_state.clone());
104 if shell_policy.allows_shell() {
105 builder.with_shell_tools().with_runtime_task_shell_tools()
106 } else {
107 builder
108 }
109 };
110
111 builder = builder
112 .with_review_tool(client, model.to_string())
113 .with_user_input_tool();
114
115 if self.config.features.enabled(Feature::WebSearch) {
116 builder = builder.with_web_tools();
117 }
118
119 // Register the `remember` tool only when the user has opted in to
120 // user-memory (#489). Without that opt-in the tool would always
121 // fail; surfacing it would just waste catalog slots.
122 if should_register_remember_tool(self.config.memory_enabled) {
123 builder = builder.with_remember_tool();
124 }
125
126 // Register image_analyze tool when vision_model is configured and feature enabled.
127 if self.config.features.enabled(Feature::VisionModel)
128 && let Some(ref vision_config) = self.config.vision_config
129 {
130 builder = builder.with_vision_tools(vision_config.clone());
131 }
132
133 // Register the `notify` tool unconditionally (#1322). It has no
134 // side effects beyond a single terminal escape write and respects
135 // the user's `[notifications].method` config (including `off`),
136 // so there's no failure mode worth gating on.
137 builder = builder.with_notify_tool();
138
139 // Register the `registry_sync` tool for fetching and caching
140 // MCP Registry server metadata. Rides on `Feature::Mcp` — the same
141 // flag that gates the rest of the MCP system (defaults to enabled;
142 // opt out via `[features]` in config.toml).
143 if self.config.features.enabled(Feature::Mcp) {
144 builder = builder.with_registry_mcp_sync_tool();
145 }
146
147 // Register the start_mcp_server tool so LLM can dynamically start
148 // MCP servers from conversation context. Only when the pool has been
149 // initialized (lazy via ensure_mcp_pool).
150 if let Some(ref pool) = self.mcp_pool {
151 builder = builder
152 .with_runtime_mcp_tool(Arc::clone(pool))
153 .with_registry_mcp_start_tool(Arc::clone(pool));
154 }
155
156 builder
157 }
158 }
159
160 #[cfg(test)]
161 mod tests {
162 use super::should_register_remember_tool;
163
164 #[test]
165 fn remember_tool_registration_requires_memory_opt_in() {
166 assert!(should_register_remember_tool(true));
167 assert!(!should_register_remember_tool(false));
168 }
169 }
170
170 lines RUST