feat: show users when quota is breached (#687)

* feat: show users when quota is breached

* fix: converts print function to use Report

* fix: removes redundant ok

* fix: lints and fmt

Author: dingfeli

crates/fig_api_client/src/clients/streaming_client.rs

@@ -119,14 +119,23 @@ impl StreamingClient {
                     )
                     .build()
                     .expect("building conversation_state should not fail");
+                let response = client
+                    .generate_assistant_response()
+                    .conversation_state(conversation_state)
+                    .send()
+                    .await;
 
-                Ok(SendMessageOutput::Codewhisperer(
-                    client
-                        .generate_assistant_response()
-                        .conversation_state(conversation_state)
-                        .send()
-                        .await?,
-                ))
+                match response {
+                    Ok(resp) => Ok(SendMessageOutput::Codewhisperer(resp)),
+                    Err(e) => {
+                        let is_quota_breach = e.raw_response().is_some_and(|resp| resp.status().as_u16() == 429);
+                        if is_quota_breach {
+                            Err(Error::QuotaBreach("quota has reached its limit"))
+                        } else {
+                            Err(e.into())
+                        }
+                    },
+                }
             },
             inner::Inner::QDeveloper(client) => {
                 let conversation_state_builder = amzn_qdeveloper_streaming_client::types::ConversationState::builder()

crates/fig_api_client/src/error.rs

@@ -44,6 +44,10 @@ pub enum Error {
     #[error("{}", SdkErrorDisplay(.0))]
     QDeveloperChatResponseStream(#[from] SdkError<QDeveloperChatResponseStreamError, RawMessage>),
 
+    // quota breach
+    #[error("quota has reached its limit")]
+    QuotaBreach(&'static str),
+
     #[error(transparent)]
     SmithyBuild(#[from] aws_smithy_types::error::operation::BuildError),
 
@@ -66,7 +70,8 @@ impl Error {
             Error::CodewhispererChatResponseStream(_)
             | Error::QDeveloperChatResponseStream(_)
             | Error::SmithyBuild(_)
-            | Error::UnsupportedConsolas(_) => false,
+            | Error::UnsupportedConsolas(_)
+            | Error::QuotaBreach(_) => false,
         }
     }
 
@@ -82,7 +87,8 @@ impl Error {
             Error::CodewhispererChatResponseStream(_)
             | Error::QDeveloperChatResponseStream(_)
             | Error::SmithyBuild(_)
-            | Error::UnsupportedConsolas(_) => false,
+            | Error::UnsupportedConsolas(_)
+            | Error::QuotaBreach(_) => false,
         }
     }
 }

crates/q_cli/src/cli/chat/error.rs

@@ -0,0 +1,36 @@
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum PromptAndSendError {
+    FigClientError(fig_api_client::Error),
+    IO(std::io::Error),
+    Report(eyre::Report),
+}
+
+impl std::fmt::Display for PromptAndSendError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PromptAndSendError::FigClientError(err) => err.fmt(f),
+            PromptAndSendError::IO(err) => err.fmt(f),
+            PromptAndSendError::Report(err) => err.fmt(f),
+        }
+    }
+}
+
+impl From<fig_api_client::Error> for PromptAndSendError {
+    fn from(value: fig_api_client::Error) -> Self {
+        PromptAndSendError::FigClientError(value)
+    }
+}
+
+impl From<std::io::Error> for PromptAndSendError {
+    fn from(value: std::io::Error) -> Self {
+        PromptAndSendError::IO(value)
+    }
+}
+
+impl From<eyre::Report> for PromptAndSendError {
+    fn from(value: eyre::Report) -> Self {
+        PromptAndSendError::Report(value)
+    }
+}

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

@@ -1,4 +1,5 @@
 mod conversation_state;
+mod error;
 mod input_source;
 mod parse;
 mod parser;
@@ -28,6 +29,7 @@ use crossterm::{
     style,
     terminal,
 };
+use error::PromptAndSendError;
 use eyre::{
     Result,
     bail,
@@ -147,6 +149,29 @@ fn load_tools() -> Result<ToolConfiguration> {
     })
 }
 
+fn print_error<W: Write>(
+    output: &mut W,
+    prepend_msg: &str,
+    report: Option<eyre::Report>,
+) -> Result<(), std::io::Error> {
+    queue!(
+        output,
+        style::SetAttribute(Attribute::Bold),
+        style::SetForegroundColor(Color::Red),
+    )?;
+    if let Some(report) = report {
+        queue!(output, style::Print(format!("{}: {:?}\n", prepend_msg, report)),)?;
+    } else {
+        queue!(output, style::Print(prepend_msg), style::Print("\n"))?;
+    }
+    queue!(
+        output,
+        style::SetForegroundColor(Color::Reset),
+        style::SetAttribute(Attribute::Reset),
+    )?;
+    output.flush()
+}
+
 /// Required fields for initializing a new chat session.
 struct ChatArgs<'o, W> {
     /// The [Write] destination for printing conversation text.
@@ -250,15 +275,26 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
                                 cursor::Show
                             )?;
                         }
-                        execute!(
-                            self.output,
-                            style::SetAttribute(Attribute::Bold),
-                            style::SetForegroundColor(Color::Red),
-                            style::Print(format!("Amazon Q is having trouble responding right now: {:?}\n", e)),
-                            style::Print("Please try again later.\n\n"),
-                            style::SetForegroundColor(Color::Reset),
-                            style::SetAttribute(Attribute::Reset),
-                        )?;
+                        match e {
+                            PromptAndSendError::FigClientError(err) => {
+                                if let fig_api_client::Error::QuotaBreach(msg) = err {
+                                    print_error(self.output, msg, None)?;
+                                } else {
+                                    print_error(
+                                        self.output,
+                                        "Amazon Q is having trouble responding right now",
+                                        Some(err.into()),
+                                    )?;
+                                }
+                            },
+                            _ => {
+                                print_error(
+                                    self.output,
+                                    "Amazon Q is having trouble responding right now",
+                                    Some(e.into()),
+                                )?;
+                            },
+                        }
                         if self.conversation_state.next_message.is_none() {
                             self.conversation_state.history.pop_back();
                         }
@@ -445,7 +481,7 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
         &mut self,
         sigint_recver: &mut tokio::sync::mpsc::Receiver<()>,
         should_terminate: &Arc<AtomicBool>,
-    ) -> Result<Option<SendMessageOutput>> {
+    ) -> Result<Option<SendMessageOutput>, error::PromptAndSendError> {
         loop {
             // Tool uses that need to be executed.
             let mut queued_tools: Vec<(String, Tool)> = Vec::new();
@@ -548,7 +584,9 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
                 self.tool_use_recursions += 1;
                 let terminal_width = self.terminal_width();
                 if self.tool_use_recursions > MAX_TOOL_USE_RECURSIONS {
-                    bail!("Exceeded max tool use recursion limit: {}", MAX_TOOL_USE_RECURSIONS);
+                    return Err(
+                        eyre::eyre!("Exceeded max tool use recursion limit: {}", MAX_TOOL_USE_RECURSIONS).into(),
+                    );
                 }
 
                 for (i, (_, tool)) in queued_tools.iter().enumerate() {