fix: stream timeout handling (#732)

Author: brandonskiser

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

@@ -677,6 +677,34 @@ where
                         },
                     }
                 },
+                Err(RecvError::StreamTimeout { source, duration }) => {
+                    error!(
+                        ?source,
+                        "Encountered a stream timeout after waiting for {}s",
+                        duration.as_secs()
+                    );
+                    if self.interactive {
+                        execute!(self.output, cursor::Hide)?;
+                        self.spinner = Some(Spinner::new(Spinners::Dots, "Dividing up the work...".to_string()));
+                    }
+                    // For stream timeouts, we'll tell the model to try and split its response into
+                    // smaller chunks.
+                    self.conversation_state
+                        .push_assistant_message(AssistantResponseMessage {
+                            message_id: None,
+                            content: "Fake message - actual message took too long to generate".to_string(),
+                            tool_uses: None,
+                        });
+                    self.conversation_state.append_new_user_message(
+                        "You took too long to respond - try to split up the work into smaller steps.".to_string(),
+                    );
+                    self.send_tool_use_telemetry().await;
+                    return Ok(ChatState::HandleResponseStream(
+                        self.client
+                            .send_message(self.conversation_state.as_sendable_conversation_state())
+                            .await?,
+                    ));
+                },
                 Err(RecvError::UnexpectedToolUseEos {
                     tool_use_id,
                     name,

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

@@ -46,7 +46,23 @@ pub enum RecvError {
     Client(#[from] fig_api_client::Error),
     #[error("{0}")]
     Json(#[from] serde_json::Error),
+    /// An error was encountered while waiting for the next event in the stream after a noticeably
+    /// long wait time.
+    ///
+    /// *Context*: the client can throw an error after ~100s of waiting with no response, likely due
+    /// to an exceptionally complex tool use taking too long to generate.
+    #[error("The stream ended after {}s: {source}", .duration.as_secs())]
+    StreamTimeout {
+        source: fig_api_client::Error,
+        duration: std::time::Duration,
+    },
     /// Unexpected end of stream while receiving a tool use.
+    ///
+    /// *Context*: the stream can unexpectedly end with `Ok(None)` while waiting for an
+    /// exceptionally complex tool use. This is due to some proxy server dropping idle
+    /// connections after some timeout is reached.
+    ///
+    /// TODO: should this be removed?
     #[error("Unexpected end of stream for tool: {} with id: {}", .name, .tool_use_id)]
     UnexpectedToolUseEos {
         tool_use_id: String,
@@ -154,7 +170,7 @@ impl ResponseParser {
                     };
                     return Ok(ResponseEvent::EndStream { message });
                 },
-                Err(err) => return Err(err.into()),
+                Err(err) => return Err(err),
             }
         }
     }
@@ -221,7 +237,7 @@ impl ResponseParser {
     }
 
     /// Returns the next event in the [SendMessageOutput] without consuming it.
-    async fn peek(&mut self) -> Result<Option<&ChatResponseStream>, fig_api_client::Error> {
+    async fn peek(&mut self) -> Result<Option<&ChatResponseStream>, RecvError> {
         if self.peek.is_some() {
             return Ok(self.peek.as_ref());
         }
@@ -235,14 +251,27 @@ impl ResponseParser {
     }
 
     /// Consumes the next [SendMessageOutput] event.
-    async fn next(&mut self) -> Result<Option<ChatResponseStream>, fig_api_client::Error> {
+    async fn next(&mut self) -> Result<Option<ChatResponseStream>, RecvError> {
         if let Some(ev) = self.peek.take() {
             return Ok(Some(ev));
         }
         trace!("Attempting to recv next event");
-        let r = self.response.recv().await?;
-        trace!(?r, "Received new event");
-        Ok(r)
+        let start = std::time::Instant::now();
+        let result = self.response.recv().await;
+        let duration = std::time::Instant::now().duration_since(start);
+        match result {
+            Ok(r) => {
+                trace!(?r, "Received new event");
+                Ok(r)
+            },
+            Err(err) => {
+                if duration.as_secs() >= 59 {
+                    Err(RecvError::StreamTimeout { source: err, duration })
+                } else {
+                    Err(err.into())
+                }
+            },
+        }
     }
 }