updates tools info per try chat loop

Author: dingfeli

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

@@ -33,6 +33,7 @@ use std::process::{
     ExitCode,
 };
 use std::sync::Arc;
+use std::sync::atomic::Ordering;
 use std::time::Duration;
 use std::{
     env,
@@ -795,6 +796,16 @@ impl ChatContext {
             let ctrl_c_stream = ctrl_c();
             debug!(?chat_state, "changing to state");
 
+            // Update conversation state with new tool information
+            if self
+                .conversation_state
+                .tool_manager
+                .has_new_stuff
+                .load(Ordering::Relaxed)
+            {
+                self.conversation_state.update_state().await;
+            }
+
             let result = match chat_state {
                 ChatState::PromptUser {
                     tool_uses,

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

@@ -32,7 +32,9 @@ pub enum UpdateEventMessage {
         server_name: String,
         result: ResourceTemplatesListResult,
     },
-    DisplayTaskEnded,
+    InitStart {
+        server_name: String,
+    },
 }
 
 #[derive(Clone, Debug)]
@@ -112,6 +114,16 @@ impl Messenger for ServerMessenger {
             .map_err(|e| MessengerError::Custom(e.to_string()))?)
     }
 
+    async fn send_init_msg(&self) -> Result<(), MessengerError> {
+        Ok(self
+            .update_event_sender
+            .send(UpdateEventMessage::InitStart {
+                server_name: self.server_name.clone(),
+            })
+            .await
+            .map_err(|e| MessengerError::Custom(e.to_string()))?)
+    }
+
     fn duplicate(&self) -> Box<dyn Messenger> {
         Box::new(self.clone())
     }

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

@@ -40,7 +40,10 @@ use serde::{
 };
 use thiserror::Error;
 use tokio::signal::ctrl_c;
-use tokio::sync::Mutex;
+use tokio::sync::{
+    Mutex,
+    RwLock,
+};
 use tracing::{
     error,
     warn,
@@ -406,6 +409,8 @@ impl ToolManagerBuilder {
         let new_tool_specs_clone = new_tool_specs.clone();
         let has_new_stuff = Arc::new(AtomicBool::new(false));
         let has_new_stuff_clone = has_new_stuff.clone();
+        let pending = Arc::new(RwLock::new(HashSet::<String>::new()));
+        let pending_clone = pending.clone();
         let (mut msg_rx, messenger_builder) = ServerMessengerBuilder::new(20);
         tokio::spawn(async move {
             while let Some(msg) = msg_rx.recv().await {
@@ -415,6 +420,7 @@ impl ToolManagerBuilder {
                 // list calls.
                 match msg {
                     UpdateEventMessage::ToolsListResult { server_name, result } => {
+                        pending_clone.write().await.remove(&server_name);
                         let mut specs = result
                             .tools
                             .into_iter()
@@ -464,14 +470,16 @@ impl ToolManagerBuilder {
                         server_name: _,
                         result: _,
                     } => {},
-                    UpdateEventMessage::DisplayTaskEnded => {
-                        load_msg_sender.take();
+                    UpdateEventMessage::InitStart { server_name } => {
+                        pending_clone.write().await.insert(server_name.clone());
+                        if let Some(sender) = &load_msg_sender {
+                            let _ = sender.send(LoadingMsg::Add(server_name));
+                        }
                     },
                 }
             }
         });
         for (mut name, init_res) in pre_initialized {
-            let _ = tx.send(LoadingMsg::Add(name.clone()));
             match init_res {
                 Ok(mut client) => {
                     let messenger = messenger_builder.build_with_name(client.get_server_name().to_owned());
@@ -592,6 +600,7 @@ impl ToolManagerBuilder {
             clients,
             prompts,
             loading_display_task,
+            pending_clients: pending,
             loading_status_sender,
             new_tool_specs,
             has_new_stuff,
@@ -638,8 +647,19 @@ pub struct ToolManager {
     /// These clients are used to communicate with MCP servers.
     pub clients: HashMap<String, Arc<CustomToolClient>>,
 
+    #[allow(dead_code)]
+    /// A list of client names that are still in the process of being initialized
+    pub pending_clients: Arc<RwLock<HashSet<String>>>,
+
+    /// Flag indicating whether new tool specifications have been added since the last update.
+    /// When set to true, it signals that the tool manager needs to refresh its internal state
+    /// to incorporate newly available tools from MCP servers.
     pub has_new_stuff: Arc<AtomicBool>,
 
+    /// Storage for newly discovered tool specifications from MCP servers that haven't yet been
+    /// integrated into the main tool registry. This field holds a thread-safe reference to a map
+    /// of server names to their tool specifications and name mappings, allowing concurrent updates
+    /// from server initialization processes.
     new_tool_specs: NewToolSpecs,
 
     /// Cache for prompts collected from different servers.
@@ -1059,6 +1079,10 @@ impl ToolManager {
         );
         Ok(())
     }
+
+    pub async fn _pending_clients(&self) -> Vec<String> {
+        self.pending_clients.read().await.iter().cloned().collect::<Vec<_>>()
+    }
 }
 
 #[inline]
@@ -1150,7 +1174,7 @@ fn process_tool_specs(
             server_name, msg
         );
         if is_in_display {
-            Some(LoadingMsg::Error {
+            Some(LoadingMsg::Warn {
                 name: server_name.to_string(),
                 msg: eyre::eyre!(msg),
             })

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

@@ -92,6 +92,9 @@ impl CustomToolClient {
                 server_capabilities,
                 ..
             } => {
+                if let Some(messenger) = &client.messenger {
+                    let _ = messenger.send_init_msg().await;
+                }
                 // We'll need to first initialize. This is the handshake every client and server
                 // needs to do before proceeding to anything else
                 let cap = client.init().await?;

crates/chat-cli/src/mcp_client/messenger.rs

@@ -33,6 +33,9 @@ pub trait Messenger: std::fmt::Debug + Send + Sync + 'static {
         result: ResourceTemplatesListResult,
     ) -> Result<(), MessengerError>;
 
+    /// Signals to the orchestrator that a server has started initializing
+    async fn send_init_msg(&self) -> Result<(), MessengerError>;
+
     /// Creates a duplicate of the messenger object
     /// This function is used to create a new instance of the messenger with the same configuration
     fn duplicate(&self) -> Box<dyn Messenger>;
@@ -68,6 +71,10 @@ impl Messenger for NullMessenger {
         Ok(())
     }
 
+    async fn send_init_msg(&self) -> Result<(), MessengerError> {
+        Ok(())
+    }
+
     fn duplicate(&self) -> Box<dyn Messenger> {
         Box::new(NullMessenger)
     }