fix: add check for tool use timeouts (#1324)

Author: brandonskiser

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

@@ -2569,16 +2569,29 @@ impl ChatContext {
                             tool_use_id,
                             name,
                             message,
+                            time_elapsed,
                         } => {
                             error!(
                                 recv_error.request_id,
                                 tool_use_id, name, "The response stream ended before the entire tool use was received"
                             );
                             if self.interactive {
-                                execute!(self.output, cursor::Hide)?;
+                                drop(self.spinner.take());
+                                queue!(
+                                    self.output,
+                                    terminal::Clear(terminal::ClearType::CurrentLine),
+                                    cursor::MoveToColumn(0),
+                                    style::SetForegroundColor(Color::Yellow),
+                                    style::SetAttribute(Attribute::Bold),
+                                    style::Print(format!(
+                                        "Warning: received an unexpected error from the model after {:.2}s\n\n",
+                                        time_elapsed.as_secs_f64()
+                                    )),
+                                    style::SetAttribute(Attribute::Reset),
+                                )?;
                                 self.spinner = Some(Spinner::new(
                                     Spinners::Dots,
-                                    "The generated tool use was too large, trying to divide up the work...".to_string(),
+                                    "Trying to divide up the work...".to_string(),
                                 ));
                             }
 

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

@@ -1,4 +1,7 @@
-use std::time::Instant;
+use std::time::{
+    Duration,
+    Instant,
+};
 
 use eyre::Result;
 use fig_api_client::clients::SendMessageOutput;
@@ -66,6 +69,7 @@ pub enum RecvErrorKind {
         tool_use_id: String,
         name: String,
         message: Box<AssistantMessage>,
+        time_elapsed: Duration,
     },
 }
 
@@ -180,29 +184,29 @@ impl ResponseParser {
     /// The arguments are the fields from the first [ChatResponseStream::ToolUseEvent] consumed.
     async fn parse_tool_use(&mut self, id: String, name: String) -> Result<AssistantToolUse, RecvError> {
         let mut tool_string = String::new();
-        let mut stop_seen = false;
         let start = Instant::now();
         while let Some(ChatResponseStream::ToolUseEvent { .. }) = self.peek().await? {
             if let Some(ChatResponseStream::ToolUseEvent { input, stop, .. }) = self.next().await? {
                 if let Some(i) = input {
                     tool_string.push_str(&i);
                 }
                 if let Some(true) = stop {
-                    stop_seen = true;
                     break;
                 }
             }
         }
         let args = match serde_json::from_str(&tool_string) {
             Ok(args) => args,
             Err(err) => {
-                // If the stream ended before we saw the final tool use event (and thus failed
-                // deserializing the tool use), this is most likely due to the backend dropping the
-                // connection. The tool was too large!
-                if self.peek().await?.is_none() && !stop_seen {
+                // If we failed deserializing after waiting for a long time, then this is most
+                // likely bedrock responding with a stop event for some reason without actually
+                // including the tool contents. Essentially, the tool was too large.
+                // Timeouts have been seen as short as ~1 minute, so setting the time to 30.
+                let time_elapsed = start.elapsed();
+                if self.peek().await?.is_none() && time_elapsed > Duration::from_secs(30) {
                     error!(
                         "Received an unexpected end of stream after spending ~{}s receiving tool events",
-                        Instant::now().duration_since(start).as_secs_f64()
+                        time_elapsed.as_secs_f64()
                     );
                     self.tool_uses.push(AssistantToolUse {
                         id: id.clone(),
@@ -211,18 +215,14 @@ impl ResponseParser {
                             [(
                                 "key".to_string(),
                                 serde_json::Value::String(
-                                    "fake tool use args - actual tool use was too large to include".to_string(),
+                                    "WARNING: the actual tool use arguments were too complicated to be generated"
+                                        .to_string(),
                                 ),
                             )]
                             .into_iter()
                             .collect(),
                         ),
                     });
-                    // let message = Box::new(AssistantMessage {
-                    //     message_id: Some(self.message_id.clone()),
-                    //     content: std::mem::take(&mut self.assistant_text),
-                    //     tool_uses: Some(self.tool_uses.clone().into_iter().map(Into::into).collect()),
-                    // });
                     let message = Box::new(AssistantMessage::new_tool_use(
                         Some(self.message_id.clone()),
                         std::mem::take(&mut self.assistant_text),
@@ -232,6 +232,7 @@ impl ResponseParser {
                         tool_use_id: id,
                         name,
                         message,
+                        time_elapsed,
                     }));
                 } else {
                     return Err(self.error(err));