moves tool manager to conversation state

Author: dingfeli

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

@@ -37,6 +37,7 @@ use super::token_counter::{
     CharCount,
     CharCounter,
 };
+use super::tool_manager::ToolManager;
 use super::tools::{
     InputSchema,
     QueuedTool,
@@ -85,6 +86,8 @@ pub struct ConversationState {
     pub tools: HashMap<ToolOrigin, Vec<Tool>>,
     /// Context manager for handling sticky context files
     pub context_manager: Option<ContextManager>,
+    /// Tool manager for handling tool and mcp related activities
+    pub tool_manager: ToolManager,
     /// Cached value representing the length of the user context message.
     context_message_length: Option<usize>,
     /// Stores the latest conversation summary created by /compact
@@ -99,6 +102,7 @@ impl ConversationState {
         tool_config: HashMap<String, ToolSpec>,
         profile: Option<String>,
         updates: Option<SharedWriter>,
+        tool_manager: ToolManager,
     ) -> Self {
         // Initialize context manager
         let context_manager = match ContextManager::new(ctx, None).await {
@@ -137,6 +141,7 @@ impl ConversationState {
                     acc
                 }),
             context_manager,
+            tool_manager,
             context_message_length: None,
             latest_summary: None,
             updates,
@@ -926,6 +931,7 @@ mod tests {
             tool_manager.load_tools().await.unwrap(),
             None,
             None,
+            tool_manager,
         )
         .await;
 
@@ -944,12 +950,14 @@ mod tests {
     async fn test_conversation_state_history_handling_with_tool_results() {
         // Build a long conversation history of tool use results.
         let mut tool_manager = ToolManager::default();
+        let tool_config = tool_manager.load_tools().await.unwrap();
         let mut conversation_state = ConversationState::new(
             Context::new(),
             "fake_conv_id",
-            tool_manager.load_tools().await.unwrap(),
+            tool_config.clone(),
             None,
             None,
+            tool_manager.clone(),
         )
         .await;
         conversation_state.set_next_user_message("start".to_string()).await;
@@ -975,9 +983,10 @@ mod tests {
         let mut conversation_state = ConversationState::new(
             Context::new(),
             "fake_conv_id",
-            tool_manager.load_tools().await.unwrap(),
+            tool_config.clone(),
             None,
             None,
+            tool_manager.clone(),
         )
         .await;
         conversation_state.set_next_user_message("start".to_string()).await;
@@ -1016,6 +1025,7 @@ mod tests {
             tool_manager.load_tools().await.unwrap(),
             None,
             None,
+            tool_manager,
         )
         .await;
 
@@ -1081,6 +1091,7 @@ mod tests {
             tool_manager.load_tools().await.unwrap(),
             None,
             Some(SharedWriter::stdout()),
+            tool_manager,
         )
         .await;
 

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

@@ -505,8 +505,6 @@ pub struct ChatContext {
     tool_use_telemetry_events: HashMap<String, ToolUseEventBuilder>,
     /// State used to keep track of tool use relation
     tool_use_status: ToolUseStatus,
-    /// Abstraction that consolidates custom tools with native ones
-    tool_manager: ToolManager,
     /// Any failed requests that could be useful for error report/debugging
     failed_request_ids: Vec<String>,
     /// Pending prompts to be sent
@@ -533,8 +531,15 @@ impl ChatContext {
     ) -> Result<Self> {
         let ctx_clone = Arc::clone(&ctx);
         let output_clone = output.clone();
-        let conversation_state =
-            ConversationState::new(ctx_clone, conversation_id, tool_config, profile, Some(output_clone)).await;
+        let conversation_state = ConversationState::new(
+            ctx_clone,
+            conversation_id,
+            tool_config,
+            profile,
+            Some(output_clone),
+            tool_manager,
+        )
+        .await;
         Ok(Self {
             ctx,
             settings,
@@ -550,7 +555,6 @@ impl ChatContext {
             conversation_state,
             tool_use_telemetry_events: HashMap::new(),
             tool_use_status: ToolUseStatus::Idle,
-            tool_manager,
             failed_request_ids: Vec::new(),
             pending_prompts: VecDeque::new(),
         })
@@ -1217,6 +1221,7 @@ impl ChatContext {
         #[cfg(unix)]
         if let Some(ref context_manager) = self.conversation_state.context_manager {
             let tool_names = self
+                .conversation_state
                 .tool_manager
                 .tn_map
                 .keys()
@@ -2152,9 +2157,10 @@ impl ChatContext {
 
                 match subcommand {
                     Some(ToolsSubcommand::Schema) => {
-                        let schema_json = serde_json::to_string_pretty(&self.tool_manager.schema).map_err(|e| {
-                            ChatError::Custom(format!("Error converting tool schema to string: {e}").into())
-                        })?;
+                        let schema_json = serde_json::to_string_pretty(&self.conversation_state.tool_manager.schema)
+                            .map_err(|e| {
+                                ChatError::Custom(format!("Error converting tool schema to string: {e}").into())
+                            })?;
                         queue!(self.output, style::Print(schema_json), style::Print("\n"))?;
                     },
                     Some(ToolsSubcommand::Trust { tool_names }) => {
@@ -2368,7 +2374,7 @@ impl ChatContext {
                     },
                     Some(PromptsSubcommand::Get { mut get_command }) => {
                         let orig_input = get_command.orig_input.take();
-                        let prompts = match self.tool_manager.get_prompt(get_command).await {
+                        let prompts = match self.conversation_state.tool_manager.get_prompt(get_command).await {
                             Ok(resp) => resp,
                             Err(e) => {
                                 match e {
@@ -2455,12 +2461,12 @@ impl ChatContext {
                             _ => None,
                         };
                         let terminal_width = self.terminal_width();
-                        let mut prompts_wl = self.tool_manager.prompts.write().map_err(|e| {
+                        let mut prompts_wl = self.conversation_state.tool_manager.prompts.write().map_err(|e| {
                             ChatError::Custom(
                                 format!("Poison error encountered while retrieving prompts: {}", e).into(),
                             )
                         })?;
-                        self.tool_manager.refresh_prompts(&mut prompts_wl)?;
+                        self.conversation_state.tool_manager.refresh_prompts(&mut prompts_wl)?;
                         let mut longest_name = "";
                         let arg_pos = {
                             let optimal_case = UnicodeWidthStr::width(longest_name) + terminal_width / 4;
@@ -3126,7 +3132,7 @@ impl ChatContext {
                 .set_tool_use_id(tool_use_id.clone())
                 .set_tool_name(tool_use.name.clone())
                 .utterance_id(self.conversation_state.message_id().map(|s| s.to_string()));
-            match self.tool_manager.get_tool_from_tool_use(tool_use) {
+            match self.conversation_state.tool_manager.get_tool_from_tool_use(tool_use) {
                 Ok(mut tool) => {
                     // Apply non-Q-generated context to tools
                     self.contextualize_tool(&mut tool);

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

@@ -13,6 +13,7 @@ use crate::mcp_client::{
     ToolsListResult,
 };
 
+#[allow(dead_code)]
 #[derive(Clone, Debug)]
 pub enum UpdateEventMessage {
     ToolsListResult {

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

@@ -596,7 +596,7 @@ enum OutOfSpecName {
 
 type NewToolSpecs = Arc<Mutex<HashMap<String, (HashMap<String, String>, Vec<ToolSpec>)>>>;
 
-#[derive(Default)]
+#[derive(Default, Debug)]
 /// Manages the lifecycle and interactions with tools from various sources, including MCP servers.
 /// This struct is responsible for initializing tools, handling tool requests, and maintaining
 /// a cache of available prompts from connected servers.
@@ -639,6 +639,21 @@ pub struct ToolManager {
     pub schema: HashMap<String, ToolSpec>,
 }
 
+impl Clone for ToolManager {
+    fn clone(&self) -> Self {
+        Self {
+            conversation_id: self.conversation_id.clone(),
+            clients: self.clients.clone(),
+            has_new_stuff: self.has_new_stuff.clone(),
+            new_tool_specs: self.new_tool_specs.clone(),
+            prompts: self.prompts.clone(),
+            tn_map: self.tn_map.clone(),
+            schema: self.schema.clone(),
+            ..Default::default()
+        }
+    }
+}
+
 impl ToolManager {
     pub async fn load_tools(&mut self) -> eyre::Result<HashMap<String, ToolSpec>> {
         let tx = self.loading_status_sender.take();

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

@@ -15,10 +15,7 @@ use serde::{
 use tokio::sync::RwLock;
 use tracing::warn;
 
-use super::{
-    InvokeOutput,
-    ToolSpec,
-};
+use super::InvokeOutput;
 use crate::cli::chat::CONTINUATION_LINE;
 use crate::cli::chat::token_counter::TokenCounter;
 use crate::mcp_client::{
@@ -88,32 +85,20 @@ impl CustomToolClient {
         })
     }
 
-    pub async fn init(&self) -> Result<(String, Vec<ToolSpec>)> {
+    pub async fn init(&self) -> Result<()> {
         match self {
             CustomToolClient::Stdio {
                 client,
-                server_name,
                 server_capabilities,
+                ..
             } => {
                 // 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?;
                 // We'll be scrapping this for background server load: https://github.com/aws/amazon-q-developer-cli/issues/1466
                 // So don't worry about the tidiness for now
-                let is_tool_supported = cap.tools.is_some();
                 server_capabilities.write().await.replace(cap);
-                // Assuming a shape of return as per https://spec.modelcontextprotocol.io/specification/2024-11-05/server/tools/#listing-tools
-                // let tools = if is_tool_supported {
-                //     // And now we make the server tell us what tools they have
-                //     let resp = client.request("tools/list", None).await?;
-                //     match resp.result.and_then(|r| r.get("tools").cloned()) {
-                //         Some(value) => serde_json::from_value::<Vec<ToolSpec>>(value)?,
-                //         None => Default::default(),
-                //     }
-                // } else {
-                //     Default::default()
-                // };
-                Ok((server_name.clone(), vec![]))
+                Ok(())
             },
         }
     }

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

@@ -82,6 +82,7 @@ pub struct ClientConfig {
     pub env: Option<HashMap<String, String>>,
 }
 
+#[allow(dead_code)]
 #[derive(Debug, Error)]
 pub enum ClientError {
     #[error(transparent)]
@@ -548,10 +549,6 @@ where
         )
     }
 
-    pub async fn shutdown(&self) -> Result<(), ClientError> {
-        Ok(self.transport.shutdown().await?)
-    }
-
     fn get_id(&self) -> u64 {
         self.current_id.fetch_add(1, Ordering::SeqCst)
     }

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

@@ -5,6 +5,7 @@ use serde::{
 use thiserror::Error;
 
 /// https://spec.modelcontextprotocol.io/specification/2024-11-05/server/utilities/pagination/#operations-supporting-pagination
+#[allow(clippy::enum_variant_names)]
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum PaginationSupportedOps {
     ResourcesList,

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

@@ -11,6 +11,7 @@ use super::{
 /// consumer. It is through this interface secondary information (i.e. information that are needed
 /// to make requests to mcp servers) are obtained passively. Consumers of client can of course
 /// choose to "actively" retrieve these information via explicitly making these requests.
+#[allow(dead_code)]
 #[async_trait::async_trait]
 pub trait Messenger: std::fmt::Debug + Send + Sync + 'static {
     /// Sends the result of a tools list operation to the consumer

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

@@ -1,3 +1,4 @@
+#![allow(dead_code)]
 use std::collections::HashMap;
 use std::sync::atomic::{
     AtomicBool,

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

@@ -31,6 +31,7 @@ impl From<std::io::Error> for TransportError {
     }
 }
 
+#[allow(dead_code)]
 #[async_trait::async_trait]
 pub trait Transport: Send + Sync + Debug + 'static {
     /// Sends a message over the transport layer.