agent: Prevent use of disabled tools (#33392)

Oleksiy Syvokon created 5 months ago

The agent now checks if a tool is enabled in the current profile before
calling it. Previously, the agent could still call disabled tools, which
commonly happened after switching profiles in the middle of a thread.

Release Notes:

- Fixed a bug where the agent could use disabled tools sometimes

Change summary

crates/agent/src/agent_profile.rs |  8 +++
crates/agent/src/thread.rs        | 78 ++++++++++++++++++--------------
2 files changed, 52 insertions(+), 34 deletions(-)

Detailed changes

crates/agent/src/agent_profile.rs

@@ -85,6 +85,14 @@ impl AgentProfile {
             .collect()
     }
 
+    pub fn is_tool_enabled(&self, source: ToolSource, tool_name: String, cx: &App) -> bool {
+        let Some(settings) = AgentSettings::get_global(cx).profiles.get(&self.id) else {
+            return false;
+        };
+
+        return Self::is_enabled(settings, source, tool_name);
+    }
+
     fn is_enabled(settings: &AgentProfileSettings, source: ToolSource, name: String) -> bool {
         match source {
             ToolSource::Native => *settings.tools.get(name.as_str()).unwrap_or(&false),

crates/agent/src/thread.rs

@@ -1770,7 +1770,7 @@ impl Thread {
                     match result.as_ref() {
                         Ok(stop_reason) => match stop_reason {
                             StopReason::ToolUse => {
-                                let tool_uses = thread.use_pending_tools(window, cx, model.clone());
+                                let tool_uses = thread.use_pending_tools(window, model.clone(), cx);
                                 cx.emit(ThreadEvent::UsePendingTools { tool_uses });
                             }
                             StopReason::EndTurn | StopReason::MaxTokens  => {
@@ -2120,8 +2120,8 @@ impl Thread {
     pub fn use_pending_tools(
         &mut self,
         window: Option<AnyWindowHandle>,
-        cx: &mut Context<Self>,
         model: Arc<dyn LanguageModel>,
+        cx: &mut Context<Self>,
     ) -> Vec<PendingToolUse> {
         self.auto_capture_telemetry(cx);
         let request =
@@ -2135,43 +2135,53 @@ impl Thread {
             .collect::<Vec<_>>();
 
         for tool_use in pending_tool_uses.iter() {
-            if let Some(tool) = self.tools.read(cx).tool(&tool_use.name, cx) {
-                if tool.needs_confirmation(&tool_use.input, cx)
-                    && !AgentSettings::get_global(cx).always_allow_tool_actions
-                {
-                    self.tool_use.confirm_tool_use(
-                        tool_use.id.clone(),
-                        tool_use.ui_text.clone(),
-                        tool_use.input.clone(),
-                        request.clone(),
-                        tool,
-                    );
-                    cx.emit(ThreadEvent::ToolConfirmationNeeded);
-                } else {
-                    self.run_tool(
-                        tool_use.id.clone(),
-                        tool_use.ui_text.clone(),
-                        tool_use.input.clone(),
-                        request.clone(),
-                        tool,
-                        model.clone(),
-                        window,
-                        cx,
-                    );
-                }
-            } else {
-                self.handle_hallucinated_tool_use(
-                    tool_use.id.clone(),
-                    tool_use.name.clone(),
-                    window,
-                    cx,
-                );
-            }
+            self.use_pending_tool(tool_use.clone(), request.clone(), model.clone(), window, cx);
         }
 
         pending_tool_uses
     }
 
+    fn use_pending_tool(
+        &mut self,
+        tool_use: PendingToolUse,
+        request: Arc<LanguageModelRequest>,
+        model: Arc<dyn LanguageModel>,
+        window: Option<AnyWindowHandle>,
+        cx: &mut Context<Self>,
+    ) {
+        let Some(tool) = self.tools.read(cx).tool(&tool_use.name, cx) else {
+            return self.handle_hallucinated_tool_use(tool_use.id, tool_use.name, window, cx);
+        };
+
+        if !self.profile.is_tool_enabled(tool.source(), tool.name(), cx) {
+            return self.handle_hallucinated_tool_use(tool_use.id, tool_use.name, window, cx);
+        }
+
+        if tool.needs_confirmation(&tool_use.input, cx)
+            && !AgentSettings::get_global(cx).always_allow_tool_actions
+        {
+            self.tool_use.confirm_tool_use(
+                tool_use.id,
+                tool_use.ui_text,
+                tool_use.input,
+                request,
+                tool,
+            );
+            cx.emit(ThreadEvent::ToolConfirmationNeeded);
+        } else {
+            self.run_tool(
+                tool_use.id,
+                tool_use.ui_text,
+                tool_use.input,
+                request,
+                tool,
+                model,
+                window,
+                cx,
+            );
+        }
+    }
+
     pub fn handle_hallucinated_tool_use(
         &mut self,
         tool_use_id: LanguageModelToolUseId,