feat(agent): hot swap (#2637)

* changes prompt list result to be sent over via messenger

* changes tool manager orchestrator tasks to keep prompts

* changes mpsc to broadcast

* restores prompt list functionality

* restore prompt get functionality

* adds api on tool manager to hotswap

* spawns task to send deinit msg via messenger

* adds slash command to hotswap agent

* modifies load tool wait time depending on context

* adds comments to retry logic for prompt completer

* fixes lint

* adds pid field to messenger message

* adds interactive menu for swapping agent

* fixes stale mcp load record

* documents build method on tool manager builder and refactor to make the build method smaller

Author: dingfeli

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

@@ -11,6 +11,7 @@ use crossterm::{
     execute,
     queue,
 };
+use dialoguer::Select;
 use syntect::easy::HighlightLines;
 use syntect::highlighting::{
     Style,
@@ -77,6 +78,9 @@ pub enum AgentSubcommand {
         #[arg(long, short)]
         name: String,
     },
+    /// Swap to a new agent at runtime
+    #[command(alias = "switch")]
+    Swap { name: Option<String> },
 }
 
 impl AgentSubcommand {
@@ -224,6 +228,49 @@ impl AgentSubcommand {
                     )?;
                 },
             },
+            Self::Swap { name } => {
+                if let Some(name) = name {
+                    session.conversation.swap_agent(os, &mut session.stderr, &name).await?;
+                } else {
+                    let labels = session
+                        .conversation
+                        .agents
+                        .agents
+                        .keys()
+                        .map(|name| name.as_str())
+                        .collect::<Vec<_>>();
+
+                    let name = {
+                        let idx = match Select::with_theme(&crate::util::dialoguer_theme())
+                            .with_prompt("Choose one of the following agents")
+                            .items(&labels)
+                            .default(1)
+                            .interact_on_opt(&dialoguer::console::Term::stdout())
+                        {
+                            Ok(sel) => {
+                                let _ = crossterm::execute!(
+                                    std::io::stdout(),
+                                    crossterm::style::SetForegroundColor(crossterm::style::Color::Magenta)
+                                );
+                                sel
+                            },
+                            // Ctrl‑C -> Err(Interrupted)
+                            Err(dialoguer::Error::IO(ref e)) if e.kind() == std::io::ErrorKind::Interrupted => None,
+                            Err(e) => {
+                                return Err(ChatError::Custom(
+                                    format!("Dialog has failed to make a selection {e}").into(),
+                                ));
+                            },
+                        };
+
+                        idx.and_then(|idx| labels.get(idx).cloned().map(str::to_string))
+                    };
+
+                    if let Some(name) = name {
+                        session.conversation.swap_agent(os, &mut session.stderr, &name).await?;
+                    }
+                }
+            },
         }
 
         Ok(ChatState::PromptUser {
@@ -239,6 +286,7 @@ impl AgentSubcommand {
             Self::Set { .. } => "set",
             Self::Schema => "schema",
             Self::SetDefault { .. } => "set_default",
+            Self::Swap { .. } => "swap",
         }
     }
 }

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

@@ -38,12 +38,14 @@ pub enum GetPromptError {
     MissingClient,
     #[error("Missing prompt name")]
     MissingPromptName,
-    #[error("Synchronization error: {0}")]
-    Synchronization(String),
     #[error("Missing prompt bundle")]
     MissingPromptInfo,
     #[error(transparent)]
     General(#[from] eyre::Report),
+    #[error("Incorrect response type received")]
+    IncorrectResponseType,
+    #[error("Missing channel")]
+    MissingChannel,
 }
 
 #[deny(missing_docs)]
@@ -76,10 +78,7 @@ impl PromptsArgs {
         }
 
         let terminal_width = session.terminal_width();
-        let mut prompts_wl = session.conversation.tool_manager.prompts.write().map_err(|e| {
-            ChatError::Custom(format!("Poison error encountered while retrieving prompts: {}", e).into())
-        })?;
-        session.conversation.tool_manager.refresh_prompts(&mut prompts_wl)?;
+        let prompts = session.conversation.tool_manager.list_prompts().await?;
         let mut longest_name = "";
         let arg_pos = {
             let optimal_case = UnicodeWidthStr::width(longest_name) + terminal_width / 4;
@@ -121,7 +120,7 @@ impl PromptsArgs {
             style::Print("\n"),
             style::Print(format!("{}\n", "▔".repeat(terminal_width))),
         )?;
-        let mut prompts_by_server: Vec<_> = prompts_wl
+        let mut prompts_by_server: Vec<_> = prompts
             .iter()
             .fold(
                 HashMap::<&String, Vec<&PromptBundle>>::new(),

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

@@ -699,6 +699,31 @@ impl ConversationState {
         }
         self.transcript.push_back(message);
     }
+
+    /// Swapping agent involves the following:
+    /// - Reinstantiate the context manager
+    /// - Swap agent on tool manager
+    pub async fn swap_agent(
+        &mut self,
+        os: &mut Os,
+        output: &mut impl Write,
+        agent_name: &str,
+    ) -> Result<(), ChatError> {
+        let agent = self.agents.switch(agent_name).map_err(ChatError::AgentSwapError)?;
+        self.context_manager.replace({
+            ContextManager::from_agent(agent, calc_max_context_files_size(self.model_info.as_ref()))
+                .map_err(|e| ChatError::Custom(format!("Context manager has failed to instantiate: {e}").into()))?
+        });
+
+        self.tool_manager
+            .swap_agent(os, output, agent)
+            .await
+            .map_err(ChatError::AgentSwapError)?;
+
+        self.update_state(true).await;
+
+        Ok(())
+    }
 }
 
 /// Represents a conversation state that can be converted into a [FigConversationState] (the type

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

@@ -1,7 +1,11 @@
 use eyre::Result;
 use rustyline::error::ReadlineError;
 
-use super::prompt::rl;
+use super::prompt::{
+    PromptQueryResponseReceiver,
+    PromptQuerySender,
+    rl,
+};
 #[cfg(unix)]
 use super::skim_integration::SkimCommandSelector;
 use crate::os::Os;
@@ -28,11 +32,7 @@ mod inner {
 }
 
 impl InputSource {
-    pub fn new(
-        os: &Os,
-        sender: std::sync::mpsc::Sender<Option<String>>,
-        receiver: std::sync::mpsc::Receiver<Vec<String>>,
-    ) -> Result<Self> {
+    pub fn new(os: &Os, sender: PromptQuerySender, receiver: PromptQueryResponseReceiver) -> Result<Self> {
         Ok(Self(inner::Inner::Readline(rl(os, sender, receiver)?)))
     }
 

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

@@ -96,6 +96,8 @@ use tokio::sync::{
     broadcast,
 };
 use tool_manager::{
+    PromptQuery,
+    PromptQueryResult,
     ToolManager,
     ToolManagerBuilder,
 };
@@ -334,11 +336,14 @@ impl ChatArgs {
             Some(default_model_opt.model_id.clone())
         };
 
-        let (prompt_request_sender, prompt_request_receiver) = std::sync::mpsc::channel::<Option<String>>();
-        let (prompt_response_sender, prompt_response_receiver) = std::sync::mpsc::channel::<Vec<String>>();
+        let (prompt_request_sender, prompt_request_receiver) = tokio::sync::broadcast::channel::<PromptQuery>(5);
+        let (prompt_response_sender, prompt_response_receiver) =
+            tokio::sync::broadcast::channel::<PromptQueryResult>(5);
         let mut tool_manager = ToolManagerBuilder::default()
-            .prompt_list_sender(prompt_response_sender)
-            .prompt_list_receiver(prompt_request_receiver)
+            .prompt_query_result_sender(prompt_response_sender)
+            .prompt_query_receiver(prompt_request_receiver)
+            .prompt_query_sender(prompt_request_sender.clone())
+            .prompt_query_result_receiver(prompt_response_receiver.resubscribe())
             .conversation_id(&conversation_id)
             .agent(agents.get_active().cloned().unwrap_or_default())
             .build(os, Box::new(std::io::stderr()), !self.no_interactive)
@@ -470,6 +475,8 @@ pub enum ChatError {
     NonInteractiveToolApproval,
     #[error("The conversation history is too large to compact")]
     CompactHistoryFailure,
+    #[error("Failed to swap to agent: {0}")]
+    AgentSwapError(eyre::Report),
 }
 
 impl ChatError {
@@ -486,6 +493,7 @@ impl ChatError {
             ChatError::GetPromptError(_) => None,
             ChatError::NonInteractiveToolApproval => None,
             ChatError::CompactHistoryFailure => None,
+            ChatError::AgentSwapError(_) => None,
         }
     }
 }
@@ -504,6 +512,7 @@ impl ReasonCode for ChatError {
             ChatError::Auth(_) => "AuthError".to_string(),
             ChatError::NonInteractiveToolApproval => "NonInteractiveToolApproval".to_string(),
             ChatError::CompactHistoryFailure => "CompactHistoryFailure".to_string(),
+            ChatError::AgentSwapError(_) => "AgentSwapError".to_string(),
         }
     }
 }
@@ -1602,6 +1611,7 @@ impl ChatSession {
                                 .await;
 
                             if matches!(chat_state, ChatState::Exit)
+                                || matches!(chat_state, ChatState::HandleResponseStream(_))
                                 || matches!(chat_state, ChatState::HandleInput { input: _ })
                                 // TODO(bskiser): this is just a hotfix for handling state changes
                                 // from manually running /compact, without impacting behavior of