Fix: Allow --trust-tools to work on MCP servers (#206)

Author: swapneils

crates/chat-cli/src/cli/chat/mod.rs

@@ -265,6 +265,15 @@ impl ChatArgs {
                 tool_permissions.trust_tool(&tool.name);
             }
         } else if let Some(trusted) = self.trust_tools.map(|vec| vec.into_iter().collect::<HashSet<_>>()) {
+            // --trust-all-tools takes precedence over --trust-tools=...
+            for tool_name in &trusted {
+                if !tool_name.is_empty() {
+                    // Store the original trust settings for later use with MCP tools
+                    tool_permissions.add_pending_trust_tool(tool_name.clone());
+                }
+            }
+
+            // Apply to currently known tools
             for tool in tool_config.values() {
                 if trusted.contains(&tool.name) {
                     tool_permissions.trust_tool(&tool.name);
@@ -1213,7 +1222,8 @@ impl ChatSession {
             style::SetForegroundColor(Color::Reset),
             style::SetAttribute(Attribute::Reset)
         )?;
-        let user_input = match self.read_user_input(&self.generate_tool_trust_prompt(), false) {
+        let prompt = self.generate_tool_trust_prompt();
+        let user_input = match self.read_user_input(&prompt, false) {
             Some(input) => input,
             None => return Ok(ChatState::Exit),
         };
@@ -1977,8 +1987,10 @@ impl ChatSession {
     }
 
     /// Helper function to generate a prompt based on the current context
-    fn generate_tool_trust_prompt(&self) -> String {
-        prompt::generate_prompt(self.conversation.current_profile(), self.all_tools_trusted())
+    fn generate_tool_trust_prompt(&mut self) -> String {
+        let profile = self.conversation.current_profile().map(|s| s.to_string());
+        let all_trusted = self.all_tools_trusted();
+        prompt::generate_prompt(profile.as_deref(), all_trusted)
     }
 
     async fn send_tool_use_telemetry(&mut self, telemetry: &TelemetryThread) {
@@ -1997,7 +2009,7 @@ impl ChatSession {
         (self.terminal_width_provider)().unwrap_or(80)
     }
 
-    fn all_tools_trusted(&self) -> bool {
+    fn all_tools_trusted(&mut self) -> bool {
         self.conversation.tools.values().flatten().all(|t| match t {
             FigTool::ToolSpecification(t) => self.tool_permissions.is_trusted(&t.name),
         })

crates/chat-cli/src/cli/chat/tools/mod.rs

@@ -7,7 +7,10 @@ pub mod knowledge;
 pub mod thinking;
 pub mod use_aws;
 
-use std::collections::HashMap;
+use std::collections::{
+    HashMap,
+    HashSet,
+};
 use std::io::Write;
 use std::path::{
     Path,
@@ -136,30 +139,38 @@ pub struct ToolPermissions {
     // We need this field for any stragglers
     pub trust_all: bool,
     pub permissions: HashMap<String, ToolPermission>,
+    // Store pending trust-tool patterns for MCP tools that may be loaded later
+    pub pending_trusted_tools: HashSet<String>,
 }
 
 impl ToolPermissions {
     pub fn new(capacity: usize) -> Self {
         Self {
             trust_all: false,
             permissions: HashMap::with_capacity(capacity),
+            pending_trusted_tools: HashSet::new(),
         }
     }
 
-    pub fn is_trusted(&self, tool_name: &str) -> bool {
+    pub fn is_trusted(&mut self, tool_name: &str) -> bool {
+        // Check if we should trust from pending patterns first
+        if self.should_trust_from_pending(tool_name) {
+            self.trust_tool(tool_name);
+            self.pending_trusted_tools.remove(tool_name);
+        }
+
         self.trust_all || self.permissions.get(tool_name).is_some_and(|perm| perm.trusted)
     }
 
     /// Returns a label to describe the permission status for a given tool.
-    pub fn display_label(&self, tool_name: &str) -> String {
-        if self.has(tool_name) || self.trust_all {
-            if self.is_trusted(tool_name) {
-                format!("  {}", "trusted".dark_green().bold())
-            } else {
-                format!("  {}", "not trusted".dark_grey())
-            }
-        } else {
-            self.default_permission_label(tool_name)
+    pub fn display_label(&mut self, tool_name: &str) -> String {
+        let is_trusted = self.is_trusted(tool_name);
+        let has_setting = self.has(tool_name) || self.trust_all;
+
+        match (has_setting, is_trusted) {
+            (true, true) => format!("  {}", "trusted".dark_green().bold()),
+            (true, false) => format!("  {}", "not trusted".dark_grey()),
+            _ => self.default_permission_label(tool_name),
         }
     }
 
@@ -170,21 +181,41 @@ impl ToolPermissions {
 
     pub fn untrust_tool(&mut self, tool_name: &str) {
         self.trust_all = false;
+        self.pending_trusted_tools.remove(tool_name);
         self.permissions
             .insert(tool_name.to_string(), ToolPermission { trusted: false });
     }
 
     pub fn reset(&mut self) {
         self.trust_all = false;
         self.permissions.clear();
+        self.pending_trusted_tools.clear();
     }
 
     pub fn reset_tool(&mut self, tool_name: &str) {
         self.trust_all = false;
         self.permissions.remove(tool_name);
+        self.pending_trusted_tools.remove(tool_name);
     }
 
-    pub fn has(&self, tool_name: &str) -> bool {
+    /// Add a pending trust pattern for tools that may be loaded later
+    pub fn add_pending_trust_tool(&mut self, pattern: String) {
+        self.pending_trusted_tools.insert(pattern);
+    }
+
+    /// Check if a tool should be trusted based on preceding trust declarations
+    pub fn should_trust_from_pending(&self, tool_name: &str) -> bool {
+        // Check for exact match
+        self.pending_trusted_tools.contains(tool_name)
+    }
+
+    pub fn has(&mut self, tool_name: &str) -> bool {
+        // Check if we should trust from pending tools first
+        if self.should_trust_from_pending(tool_name) {
+            self.trust_tool(tool_name);
+            self.pending_trusted_tools.remove(tool_name);
+        }
+
         self.permissions.contains_key(tool_name)
     }