feat: Prompt users to select a different model when at capacity (#451)

Co-authored-by: YIFAN LIU <[email protected]>
Co-authored-by: Brandon Kiser <[email protected]>

Author: 3 people

Cargo.toml

@@ -9,6 +9,7 @@ members = [
     "tests/fig-api/fig-api-mock",
     "tests/figterm2",
 ]
+default-members = ["crates/chat-cli"]
 
 [workspace.package]
 authors = [

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

@@ -46,60 +46,66 @@ pub struct ModelArgs;
 
 impl ModelArgs {
     pub async fn execute(self, session: &mut ChatSession) -> Result<ChatState, ChatError> {
-        queue!(session.stderr, style::Print("\n"))?;
-        let active_model_id = session.conversation.model.as_deref();
-        let labels: Vec<String> = MODEL_OPTIONS
-            .iter()
-            .map(|opt| {
-                if (opt.model_id.is_empty() && active_model_id.is_none()) || Some(opt.model_id) == active_model_id {
-                    format!("{} (active)", opt.name)
-                } else {
-                    opt.name.to_owned()
-                }
-            })
-            .collect();
-
-        let selection: Option<_> = match Select::with_theme(&crate::util::dialoguer_theme())
-            .with_prompt("Select a model for this chat session")
-            .items(&labels)
-            .default(0)
-            .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!("Failed to choose model: {e}").into())),
-        };
+        Ok(select_model(session)?.unwrap_or(ChatState::PromptUser {
+            skip_printing_tools: false,
+        }))
+    }
+}
 
-        queue!(session.stderr, style::ResetColor)?;
+pub fn select_model(session: &mut ChatSession) -> Result<Option<ChatState>, ChatError> {
+    queue!(session.stderr, style::Print("\n"))?;
+    let active_model_id = session.conversation.model.as_deref();
+    let labels: Vec<String> = MODEL_OPTIONS
+        .iter()
+        .map(|opt| {
+            if (opt.model_id.is_empty() && active_model_id.is_none()) || Some(opt.model_id) == active_model_id {
+                format!("{} (active)", opt.name)
+            } else {
+                opt.name.to_owned()
+            }
+        })
+        .collect();
 
-        if let Some(index) = selection {
-            let selected = &MODEL_OPTIONS[index];
-            let model_id_str = selected.model_id.to_string();
-            session.conversation.model = Some(model_id_str);
+    let selection: Option<_> = match Select::with_theme(&crate::util::dialoguer_theme())
+        .with_prompt("Select a model for this chat session")
+        .items(&labels)
+        .default(0)
+        .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 => return Ok(None),
+        Err(e) => return Err(ChatError::Custom(format!("Failed to choose model: {e}").into())),
+    };
 
-            queue!(
-                session.stderr,
-                style::Print("\n"),
-                style::Print(format!(" Using {}\n\n", selected.name)),
-                style::ResetColor,
-                style::SetForegroundColor(Color::Reset),
-                style::SetBackgroundColor(Color::Reset),
-            )?;
-        }
+    queue!(session.stderr, style::ResetColor)?;
 
-        execute!(session.stderr, style::ResetColor)?;
+    if let Some(index) = selection {
+        let selected = &MODEL_OPTIONS[index];
+        let model_id_str = selected.model_id.to_string();
+        session.conversation.model = Some(model_id_str);
 
-        Ok(ChatState::PromptUser {
-            skip_printing_tools: false,
-        })
+        queue!(
+            session.stderr,
+            style::Print("\n"),
+            style::Print(format!(" Using {}\n\n", selected.name)),
+            style::ResetColor,
+            style::SetForegroundColor(Color::Reset),
+            style::SetBackgroundColor(Color::Reset),
+        )?;
     }
+
+    execute!(session.stderr, style::ResetColor)?;
+
+    Ok(Some(ChatState::PromptUser {
+        skip_printing_tools: false,
+    }))
 }
 
 /// Returns Claude 3.7 for: Amazon IDC users, FRA region users

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

@@ -39,6 +39,7 @@ use clap::{
     Parser,
 };
 use cli::compact::CompactStrategy;
+use cli::model::select_model;
 use context::ContextManager;
 pub use conversation::ConversationState;
 use conversation::TokenWarningLevel;
@@ -660,6 +661,12 @@ impl ChatSession {
                     Err(ChatError::Interrupted { tool_uses: None })
                 }
             },
+            ChatState::RetryModelOverload => tokio::select! {
+                res = self.retry_model_overload(os) => res,
+                Ok(_) = ctrl_c_stream => {
+                    Err(ChatError::Interrupted { tool_uses: None })
+                }
+            },
             ChatState::Exit => return Ok(()),
         };
 
@@ -776,15 +783,49 @@ impl ChatSession {
                 },
                 ApiClientError::QuotaBreach { message, .. } => (message, Report::from(err), true),
                 ApiClientError::ModelOverloadedError { request_id, .. } => {
+                    if self.interactive {
+                        execute!(
+                            self.stderr,
+                            style::SetAttribute(Attribute::Bold),
+                            style::SetForegroundColor(Color::Red),
+                            style::Print(
+                                "\nThe model you've selected is temporarily unavailable. Please select a different model.\n"
+                            ),
+                            style::SetAttribute(Attribute::Reset),
+                            style::SetForegroundColor(Color::Reset),
+                        )?;
+
+                        if let Some(id) = request_id {
+                            self.conversation
+                                .append_transcript(format!("Model unavailable (Request ID: {})", id));
+                        }
+
+                        self.inner = Some(ChatState::RetryModelOverload);
+
+                        return Ok(());
+                    }
+
+                    // non-interactive throws this error
+                    let model_instruction = "Please relaunch with '--model <model_id>' to use a different model.";
                     let err = format!(
-                        "The model you've selected is temporarily unavailable. Please use '/model' to select a different model and try again.{}\n\n",
+                        "The model you've selected is temporarily unavailable. {}{}\n\n",
+                        model_instruction,
                         match request_id {
                             Some(id) => format!("\n    Request ID: {}", id),
                             None => "".to_owned(),
                         }
                     );
                     self.conversation.append_transcript(err.clone());
-                    ("Amazon Q is having trouble responding right now", eyre!(err), true)
+                    execute!(
+                        self.stderr,
+                        style::SetAttribute(Attribute::Bold),
+                        style::SetForegroundColor(Color::Red),
+                        style::Print("Amazon Q is having trouble responding right now:\n"),
+                        style::Print(format!("    {}\n", err.clone())),
+                        style::SetAttribute(Attribute::Reset),
+                        style::SetForegroundColor(Color::Reset),
+                    )?;
+                    ("Amazon Q is having trouble responding right now", eyre!(err), false)
                 },
                 ApiClientError::MonthlyLimitReached { .. } => {
                     let subscription_status = get_subscription_status(os).await;
@@ -935,6 +976,8 @@ enum ChatState {
         /// Parameters for how to perform the compaction request.
         strategy: CompactStrategy,
     },
+    /// Retry the current request if we encounter a model overloaded error.
+    RetryModelOverload,
     /// Exit the chat.
     Exit,
 }
@@ -2086,6 +2129,35 @@ impl ChatSession {
         Ok(ChatState::ExecuteTools)
     }
 
+    async fn retry_model_overload(&mut self, os: &mut Os) -> Result<ChatState, ChatError> {
+        match select_model(self) {
+            Ok(Some(_)) => (),
+            Ok(None) => {
+                // User did not select a model, so reset the current request state.
+                self.conversation.enforce_conversation_invariants();
+                self.conversation.reset_next_user_message();
+                self.pending_tool_index = None;
+                return Ok(ChatState::PromptUser {
+                    skip_printing_tools: false,
+                });
+            },
+            Err(err) => return Err(err),
+        }
+
+        let conv_state = self
+            .conversation
+            .as_sendable_conversation_state(os, &mut self.stderr, true)
+            .await?;
+
+        if self.interactive {
+            self.spinner = Some(Spinner::new(Spinners::Dots, "Thinking...".to_owned()));
+        }
+
+        Ok(ChatState::HandleResponseStream(
+            os.client.send_message(conv_state).await?,
+        ))
+    }
+
     /// Apply program context to tools that Q may not have.
     // We cannot attach this any other way because Tools are constructed by deserializing
     // output from Amazon Q.