fix: show request id on errors encountered during the response stream (#842)

Author: brandonskiser

crates/fig_api_client/src/clients/streaming_client.rs

@@ -178,6 +178,14 @@ pub enum SendMessageOutput {
 }
 
 impl SendMessageOutput {
+    pub fn request_id(&self) -> Option<&str> {
+        match self {
+            SendMessageOutput::Codewhisperer(output) => output.request_id(),
+            SendMessageOutput::QDeveloper(output) => output.request_id(),
+            SendMessageOutput::Mock(_) => None,
+        }
+    }
+
     pub async fn recv(&mut self) -> Result<Option<ChatResponseStream>, Error> {
         match self {
             SendMessageOutput::Codewhisperer(output) => Ok(output

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

@@ -51,6 +51,7 @@ use fig_util::CLI_BINARY_NAME;
 use input_source::InputSource;
 use parser::{
     RecvError,
+    RecvErrorKind,
     ResponseParser,
     ToolUse,
 };
@@ -1065,8 +1066,12 @@ where
                         },
                     }
                 },
-                Err(RecvError::StreamTimeout { source, duration }) => {
+                Err(RecvError {
+                    request_id,
+                    source: RecvErrorKind::StreamTimeout { source, duration },
+                }) => {
                     error!(
+                        request_id,
                         ?source,
                         "Encountered a stream timeout after waiting for {}s",
                         duration.as_secs()
@@ -1095,14 +1100,18 @@ where
                             .await?,
                     ));
                 },
-                Err(RecvError::UnexpectedToolUseEos {
-                    tool_use_id,
-                    name,
-                    message,
+                Err(RecvError {
+                    request_id,
+                    source:
+                        RecvErrorKind::UnexpectedToolUseEos {
+                            tool_use_id,
+                            name,
+                            message,
+                        },
                 }) => {
                     error!(
-                        tool_use_id,
-                        name, "The response stream ended before the entire tool use was received"
+                        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)?;

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

@@ -41,7 +41,26 @@ impl From<ToolUse> for FigToolUse {
 }
 
 #[derive(Debug, Error)]
-pub enum RecvError {
+pub struct RecvError {
+    /// The request id associated with the [SendMessageOutput] stream.
+    pub request_id: Option<String>,
+    #[source]
+    pub source: RecvErrorKind,
+}
+
+impl std::fmt::Display for RecvError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Failed to receive the next message: ")?;
+        if let Some(request_id) = self.request_id.as_ref() {
+            write!(f, "request_id: {}, error: ", request_id)?;
+        }
+        write!(f, "{}", self.source)?;
+        Ok(())
+    }
+}
+
+#[derive(Debug, Error)]
+pub enum RecvErrorKind {
     #[error("{0}")]
     Client(#[from] fig_api_client::Error),
     #[error("{0}")]
@@ -223,13 +242,13 @@ impl ResponseParser {
                         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 {
+                    return Err(self.error(RecvErrorKind::UnexpectedToolUseEos {
                         tool_use_id: id,
                         name,
                         message,
-                    });
+                    }));
                 } else {
-                    return Err(err.into());
+                    return Err(self.error(err));
                 }
             },
         };
@@ -266,13 +285,25 @@ impl ResponseParser {
             },
             Err(err) => {
                 if duration.as_secs() >= 59 {
-                    Err(RecvError::StreamTimeout { source: err, duration })
+                    Err(self.error(RecvErrorKind::StreamTimeout { source: err, duration }))
                 } else {
-                    Err(err.into())
+                    Err(self.error(err))
                 }
             },
         }
     }
+
+    fn request_id(&self) -> Option<&str> {
+        self.response.request_id()
+    }
+
+    /// Helper to create a new [RecvError] populated with the associated request id for the stream.
+    fn error(&self, source: impl Into<RecvErrorKind>) -> RecvError {
+        RecvError {
+            request_id: self.request_id().map(str::to_string),
+            source: source.into(),
+        }
+    }
 }
 
 #[derive(Debug)]