makes main chat loop update state if applicable

Author: dingfeli

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

@@ -3,6 +3,7 @@ use std::collections::{
     VecDeque,
 };
 use std::sync::Arc;
+use std::sync::atomic::Ordering;
 
 use crossterm::style::Color;
 use crossterm::{
@@ -354,6 +355,7 @@ impl ConversationState {
     /// - `run_hooks` - whether hooks should be executed and included as context
     pub async fn as_sendable_conversation_state(&mut self, run_hooks: bool) -> FigConversationState {
         debug_assert!(self.next_message.is_some());
+        self.update_state().await;
         self.enforce_conversation_invariants();
         self.history.drain(self.valid_history_range.1..);
         self.history.drain(..self.valid_history_range.0);
@@ -379,6 +381,30 @@ impl ConversationState {
             .expect("unable to construct conversation state")
     }
 
+    pub async fn update_state(&mut self) {
+        let needs_update = self.tool_manager.has_new_stuff.load(Ordering::Acquire);
+        if !needs_update {
+            return;
+        }
+        self.tool_manager.update().await;
+        self.tools = self
+            .tool_manager
+            .schema
+            .values()
+            .fold(HashMap::<ToolOrigin, Vec<Tool>>::new(), |mut acc, v| {
+                let tool = Tool::ToolSpecification(ToolSpecification {
+                    name: v.name.clone(),
+                    description: v.description.clone(),
+                    input_schema: v.input_schema.clone().into(),
+                });
+                acc.entry(v.tool_origin.clone())
+                    .and_modify(|tools| tools.push(tool.clone()))
+                    .or_insert(vec![tool]);
+                acc
+            });
+        self.tool_manager.has_new_stuff.store(false, Ordering::Release);
+    }
+
     /// Returns a conversation state representation which reflects the exact conversation to send
     /// back to the model.
     pub async fn backend_conversation_state(&mut self, run_hooks: bool, quiet: bool) -> BackendConversationState<'_> {

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

@@ -401,6 +401,7 @@ pub async fn chat(
     let tool_config = tool_manager.load_tools().await?;
     let mut tool_permissions = ToolPermissions::new(tool_config.len());
     if accept_all || trust_all_tools {
+        tool_permissions.trust_all = true;
         for tool in tool_config.values() {
             tool_permissions.trust_tool(&tool.name);
         }
@@ -2259,7 +2260,7 @@ impl ChatContext {
                         )?;
                     },
                     Some(ToolsSubcommand::ResetSingle { tool_name }) => {
-                        if self.tool_permissions.has(&tool_name) {
+                        if self.tool_permissions.has(&tool_name) || self.tool_permissions.trust_all {
                             self.tool_permissions.reset_tool(&tool_name);
                             queue!(
                                 self.output,
@@ -2735,11 +2736,9 @@ impl ChatContext {
             }
 
             // If there is an override, we will use it. Otherwise fall back to Tool's default.
-            let allowed = if self.tool_permissions.has(&tool.name) {
-                self.tool_permissions.is_trusted(&tool.name)
-            } else {
-                !tool.tool.requires_acceptance(&self.ctx)
-            };
+            let allowed = self.tool_permissions.trust_all
+                || (self.tool_permissions.has(&tool.name) && self.tool_permissions.is_trusted(&tool.name))
+                || !tool.tool.requires_acceptance(&self.ctx);
 
             if self.settings.get_bool_or("chat.enableNotifications", false) {
                 play_notification_bell(!allowed);

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

@@ -1,4 +1,7 @@
-use std::collections::HashMap;
+use std::collections::{
+    HashMap,
+    HashSet,
+};
 use std::future::Future;
 use std::hash::{
     DefaultHasher,
@@ -103,7 +106,7 @@ pub enum GetPromptError {
 /// Messages used for communication between the tool initialization thread and the loading
 /// display thread. These messages control the visual loading indicators shown to
 /// the user during tool initialization.
-pub enum LoadingMsg {
+enum LoadingMsg {
     /// Indicates a new tool is being initialized and should be added to the loading
     /// display. The String parameter is the name of the tool being initialized.
     Add(String),
@@ -421,7 +424,7 @@ impl ToolManagerBuilder {
                             .insert(server_name, (sanitized_mapping, specs));
                         // We only want to set this flag when the display task has ended
                         if load_msg_sender.is_none() {
-                            has_new_stuff_clone.store(true, Ordering::Relaxed);
+                            has_new_stuff_clone.store(true, Ordering::Release);
                         }
                     },
                     UpdateEventMessage::PromptsListResult {
@@ -710,28 +713,8 @@ impl ToolManager {
                 }
             }
         }
-        let new_tools = {
-            let mut new_tool_specs = self.new_tool_specs.lock().await;
-            new_tool_specs.drain().fold(HashMap::new(), |mut acc, (k, v)| {
-                acc.insert(k, v);
-                acc
-            })
-        };
-        for (_server_name, (tool_name_map, specs)) in new_tools {
-            for (k, v) in tool_name_map {
-                self.tn_map.insert(k, v);
-            }
-            for spec in specs {
-                tool_specs.insert(spec.name.clone(), spec);
-            }
-        }
-        // caching the tool names for skim operations
-        for tool_name in tool_specs.keys() {
-            if !self.tn_map.contains_key(tool_name) {
-                self.tn_map.insert(tool_name.clone(), tool_name.clone());
-            }
-        }
-        self.schema = tool_specs.clone();
+        self.update().await;
+        tool_specs.extend(self.schema.clone());
         Ok(tool_specs)
     }
 
@@ -831,6 +814,49 @@ impl ToolManager {
         })
     }
 
+    /// Updates tool managers various states with new information
+    pub async fn update(&mut self) {
+        // A hashmap of <tool name, tool spec>
+        let mut tool_specs = HashMap::<String, ToolSpec>::new();
+        let new_tools = {
+            let mut new_tool_specs = self.new_tool_specs.lock().await;
+            new_tool_specs.drain().fold(HashMap::new(), |mut acc, (k, v)| {
+                acc.insert(k, v);
+                acc
+            })
+        };
+        let mut updated_servers = HashSet::<ToolOrigin>::new();
+        for (_server_name, (tool_name_map, specs)) in new_tools {
+            // In a populated tn map (i.e. a partially initialized or outdated fleet of servers) there
+            // will be incoming tools with names that are already in the tn map, we will be writing
+            // over them (perhaps with the same information that they already had), and that's okay.
+            // In an event where a server has removed tools, the tools that are no longer available
+            // will linger in this map. This is also okay to not clean up as it does not affect the
+            // look up of tool names that are still active.
+            for (k, v) in tool_name_map {
+                self.tn_map.insert(k, v);
+            }
+            if let Some(spec) = specs.first() {
+                updated_servers.insert(spec.tool_origin.clone());
+            }
+            for spec in specs {
+                tool_specs.insert(spec.name.clone(), spec);
+            }
+        }
+        // Caching the tool names for skim operations
+        for tool_name in tool_specs.keys() {
+            if !self.tn_map.contains_key(tool_name) {
+                self.tn_map.insert(tool_name.clone(), tool_name.clone());
+            }
+        }
+        // Update schema
+        // As we are writing over the ensemble of tools in a given server, we will need to first
+        // remove everything that it has.
+        self.schema
+            .retain(|_tool_name, spec| !updated_servers.contains(&spec.tool_origin));
+        self.schema.extend(tool_specs);
+    }
+
     #[allow(clippy::await_holding_lock)]
     pub async fn get_prompt(&self, get_command: PromptsGetCommand) -> Result<JsonRpcResponse, GetPromptError> {
         let (server_name, prompt_name) = match get_command.params.name.split_once('/') {

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

@@ -126,30 +126,33 @@ pub struct ToolPermission {
 /// Tools that do not have an associated ToolPermission should use
 /// their default logic to determine to permission.
 pub struct ToolPermissions {
+    // We need this field for any stragglers
+    pub trust_all: bool,
     pub permissions: HashMap<String, ToolPermission>,
 }
 
 impl ToolPermissions {
     pub fn new(capacity: usize) -> Self {
         Self {
+            trust_all: false,
             permissions: HashMap::with_capacity(capacity),
         }
     }
 
     pub fn is_trusted(&self, tool_name: &str) -> bool {
-        self.permissions.get(tool_name).is_some_and(|perm| perm.trusted)
+        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) {
+        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)
+            self.default_permission_label(tool_name)
         }
     }
 
@@ -159,15 +162,18 @@ impl ToolPermissions {
     }
 
     pub fn untrust_tool(&mut self, tool_name: &str) {
+        self.trust_all = false;
         self.permissions
             .insert(tool_name.to_string(), ToolPermission { trusted: false });
     }
 
     pub fn reset(&mut self) {
+        self.trust_all = false;
         self.permissions.clear();
     }
 
     pub fn reset_tool(&mut self, tool_name: &str) {
+        self.trust_all = false;
         self.permissions.remove(tool_name);
     }
 
@@ -178,14 +184,15 @@ impl ToolPermissions {
     /// Provide default permission labels for the built-in set of tools.
     /// Unknown tools are assumed to be "Per-request"
     // This "static" way avoids needing to construct a tool instance.
-    fn default_permission_label(tool_name: &str) -> String {
+    fn default_permission_label(&self, tool_name: &str) -> String {
         let label = match tool_name {
             "fs_read" => "trusted".dark_green().bold(),
             "fs_write" => "not trusted".dark_grey(),
             "execute_bash" => "trust read-only commands".dark_grey(),
             "use_aws" => "trust read-only commands".dark_grey(),
             "report_issue" => "trusted".dark_green().bold(),
             "thinking" => "trusted (prerelease)".dark_green().bold(),
+            _ if self.trust_all => "trusted".dark_grey().bold(),
             _ => "not trusted".dark_grey(),
         };