fix: break up tool uses into multiple parts if stream ends unexpectedly

Author: brandonskiser

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

@@ -207,6 +207,10 @@ struct ChatContext<'o, W> {
     tool_use_recursions: u32,
     current_user_input_id: Option<String>,
     tool_use_events: Vec<ToolUseEventBuilder>,
+
+    /// Whether or not an unexpected end of chat stream was encountered while consuming the model
+    /// response's tool use data. Pair of (tool_use_id, name) for the tool that was being received.
+    encountered_tool_use_eos: Option<(String, String)>,
 }
 
 impl<'o, W> ChatContext<'o, W>
@@ -228,6 +232,7 @@ where
             tool_use_recursions: 0,
             current_user_input_id: None,
             tool_use_events: vec![],
+            encountered_tool_use_eos: None,
         }
     }
 
@@ -381,19 +386,43 @@ 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!(
-                                "We're having trouble responding right now, please try again later: {:?}",
-                                err
-                            )),
-                            style::SetForegroundColor(Color::Reset),
-                            style::SetAttribute(Attribute::Reset),
-                        )?;
-                        if self.conversation_state.next_message.is_none() {
-                            self.conversation_state.history.pop_back();
+                        match err {
+                            parser::RecvError::UnexpectedToolUseEos {
+                                tool_use_id,
+                                name,
+                                message,
+                            } => {
+                                error!(
+                                    tool_use_id,
+                                    name, "The response stream ended before the entire tool use was received"
+                                );
+                                self.conversation_state.push_assistant_message(*message);
+                                self.encountered_tool_use_eos = Some((tool_use_id, name));
+                                if self.is_interactive {
+                                    execute!(self.output, cursor::Hide)?;
+                                    self.spinner = Some(Spinner::new(
+                                        Spinners::Dots,
+                                        "The generated tool use was too large, trying to divide up the work..."
+                                            .to_string(),
+                                    ));
+                                }
+                            },
+                            err => {
+                                execute!(
+                                    self.output,
+                                    style::SetAttribute(Attribute::Bold),
+                                    style::SetForegroundColor(Color::Red),
+                                    style::Print(format!(
+                                        "We're having trouble responding right now, please try again later: {:?}",
+                                        err
+                                    )),
+                                    style::SetForegroundColor(Color::Reset),
+                                    style::SetAttribute(Attribute::Reset),
+                                )?;
+                                if self.conversation_state.next_message.is_none() {
+                                    self.conversation_state.history.pop_back();
+                                }
+                            },
                         }
                         break;
                     },
@@ -490,6 +519,23 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
         sigint_recver: &mut tokio::sync::mpsc::Receiver<()>,
         should_terminate: &Arc<AtomicBool>,
     ) -> Result<Option<SendMessageOutput>, error::PromptAndSendError> {
+        if let Some((tool_use_id, _)) = self.encountered_tool_use_eos.take() {
+            let tool_results = vec![ToolResult {
+                tool_use_id,
+                content: vec![ToolResultContentBlock::Text(
+                    "The generated tool was too large, try again but this time split up the work between multiple tool uses".to_string(),
+                )],
+                status: ToolResultStatus::Error,
+            }];
+            self.conversation_state.add_tool_results(tool_results);
+            self.send_tool_use_telemetry().await;
+            return Ok(Some(
+                self.client
+                    .send_message(self.conversation_state.as_sendable_conversation_state())
+                    .await?,
+            ));
+        }
+
         loop {
             // Tool uses that need to be executed.
             let mut queued_tools: Vec<(String, Tool)> = Vec::new();

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

@@ -1,3 +1,5 @@
+use std::time::Instant;
+
 use eyre::Result;
 use fig_api_client::clients::SendMessageOutput;
 use fig_api_client::model::{
@@ -6,6 +8,7 @@ use fig_api_client::model::{
     ChatResponseStream,
     ToolUse as FigToolUse,
 };
+use thiserror::Error;
 use tracing::{
     error,
     trace,
@@ -34,6 +37,21 @@ impl From<ToolUse> for FigToolUse {
     }
 }
 
+#[derive(Debug, Error)]
+pub enum RecvError {
+    #[error("{0}")]
+    Client(#[from] fig_api_client::Error),
+    #[error("{0}")]
+    Json(#[from] serde_json::Error),
+    /// Unexpected end of stream while receiving a tool use.
+    #[error("Unexpected end of stream for tool: {} with id: {}", .name, .tool_use_id)]
+    UnexpectedToolUseEos {
+        tool_use_id: String,
+        name: String,
+        message: Box<Message>,
+    },
+}
+
 /// State associated with parsing a [ConverseStreamResponse] into a [Message].
 ///
 /// # Usage
@@ -70,7 +88,7 @@ impl ResponseParser {
     }
 
     /// Consumes the associated [ConverseStreamResponse] until a valid [ResponseEvent] is parsed.
-    pub async fn recv(&mut self) -> Result<ResponseEvent> {
+    pub async fn recv(&mut self) -> Result<ResponseEvent, RecvError> {
         if let Some((id, name)) = self.parsing_tool_use.take() {
             let tool_use = self.parse_tool_use(id, name).await?;
             self.tool_uses.push(tool_use.clone());
@@ -150,19 +168,63 @@ impl ResponseParser {
     /// Consumes the response stream until a valid [ToolUse] is parsed.
     ///
     /// The arguments are the fields from the first [ChatResponseStream::ToolUseEvent] consumed.
-    async fn parse_tool_use(&mut self, id: String, name: String) -> Result<ToolUse> {
+    async fn parse_tool_use(&mut self, id: String, name: String) -> Result<ToolUse, 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 = serde_json::from_str(&tool_string)?;
+        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 {
+                    error!(
+                        "Received an unexpected end of stream after spending ~{}s receiving tool events",
+                        Instant::now().duration_since(start).as_secs_f64()
+                    );
+                    self.tool_uses.push(ToolUse {
+                        id: id.clone(),
+                        name: name.clone(),
+                        args: serde_json::Value::Object(
+                            [(
+                                "key".to_string(),
+                                serde_json::Value::String(
+                                    "fake tool use args - actual tool use was too large to include".to_string(),
+                                ),
+                            )]
+                            .into_iter()
+                            .collect(),
+                        ),
+                    });
+                    let message = Box::new(Message(ChatMessage::AssistantResponseMessage(
+                        AssistantResponseMessage {
+                            message_id: self.message_id.take(),
+                            content: std::mem::take(&mut self.assistant_text),
+                            tool_uses: Some(self.tool_uses.clone().into_iter().map(Into::into).collect()),
+                        },
+                    )));
+                    return Err(RecvError::UnexpectedToolUseEos {
+                        tool_use_id: id,
+                        name,
+                        message,
+                    });
+                } else {
+                    return Err(err.into());
+                }
+            },
+        };
         Ok(ToolUse { id, name, args })
     }