fix: response stream parser buffering assistant text until the end (#685)

* fix: response stream parser buffering assistant text until the end

* fix parser bug

Author: brandonskiser

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

@@ -289,7 +289,7 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
                 }
             }
 
-            let mut waiting_for_tool = false;
+            let mut tool_name_being_recvd: Option<String> = None;
             loop {
                 match parser.recv().await {
                     Ok(msg_event) => {
@@ -307,14 +307,10 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
                                 });
                             },
                             parser::ResponseEvent::ToolUseStart { name } => {
-                                if self.is_interactive {
-                                    execute!(self.output, style::Print("\n\n"))?;
-                                    self.spinner = Some(Spinner::new(
-                                        Spinners::Dots,
-                                        format!("Creating tool {}...", name.green()),
-                                    ));
-                                }
-                                waiting_for_tool = true;
+                                // We need to flush the buffer here, otherwise text will not be
+                                // printed while we are receiving tool use events.
+                                buf.push('\n');
+                                tool_name_being_recvd = Some(name);
                             },
                             parser::ResponseEvent::AssistantText(text) => {
                                 buf.push_str(&text);
@@ -330,7 +326,7 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
                                     )?;
                                 }
                                 self.tool_uses.push(tool_use);
-                                waiting_for_tool = false;
+                                tool_name_being_recvd = None;
                             },
                             parser::ResponseEvent::EndStream { message } => {
                                 self.conversation_state.push_assistant_message(message);
@@ -373,7 +369,7 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
                     buf.push('\n');
                 }
 
-                if !waiting_for_tool && !buf.is_empty() && self.is_interactive && self.spinner.is_some() {
+                if tool_name_being_recvd.is_none() && !buf.is_empty() && self.is_interactive && self.spinner.is_some() {
                     drop(self.spinner.take());
                     queue!(
                         self.output,
@@ -400,6 +396,15 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
                     }
                 }
 
+                // Set spinner after showing all of the assistant text content so far.
+                if let (Some(name), true) = (&tool_name_being_recvd, self.is_interactive) {
+                    execute!(self.output, style::Print("\n"))?;
+                    self.spinner = Some(Spinner::new(
+                        Spinners::Dots,
+                        format!("Creating tool {}...", name.clone().green()),
+                    ));
+                }
+
                 if ended {
                     if let (Some(conversation_id), Some(message_id)) = (
                         self.conversation_state.conversation_id(),

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

@@ -52,10 +52,6 @@ pub struct ResponseParser {
     assistant_text: String,
     /// Tool uses requested by the model.
     tool_uses: Vec<ToolUse>,
-    /// Buffered line required in case we need to discard a code reference event
-    buffered_line: Option<String>,
-    /// Short circuit and return early since we simply need to clear our buffered line
-    short_circuit: bool,
     /// Whether or not we are currently receiving tool use delta events. Tuple of
     /// `Some((tool_use_id, name))` if true, [None] otherwise.
     parsing_tool_use: Option<(String, String)>,
@@ -69,46 +65,40 @@ impl ResponseParser {
             message_id: None,
             assistant_text: String::new(),
             tool_uses: Vec::new(),
-            buffered_line: None,
-            short_circuit: false,
             parsing_tool_use: None,
         }
     }
 
     /// Consumes the associated [ConverseStreamResponse] until a valid [ResponseEvent] is parsed.
     pub async fn recv(&mut self) -> Result<ResponseEvent> {
-        if self.short_circuit {
-            let message = Message(ChatMessage::AssistantResponseMessage(AssistantResponseMessage {
-                message_id: self.message_id.take(),
-                content: std::mem::take(&mut self.assistant_text),
-                tool_uses: if self.tool_uses.is_empty() {
-                    None
-                } else {
-                    Some(self.tool_uses.clone().into_iter().map(Into::into).collect())
-                },
-            }));
-            return Ok(ResponseEvent::EndStream { message });
-        }
-
         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());
             return Ok(ResponseEvent::ToolUse(tool_use));
         }
 
+        // First, handle discarding AssistantResponseEvent's that immediately precede a
+        // CodeReferenceEvent.
+        let peek = self.peek().await?;
+        if let Some(ChatResponseStream::AssistantResponseEvent { content }) = peek {
+            // Cloning to bypass borrowchecker stuff.
+            let content = content.clone();
+            self.next().await?;
+            match self.peek().await? {
+                Some(ChatResponseStream::CodeReferenceEvent(_)) => (),
+                _ => {
+                    self.assistant_text.push_str(&content);
+                    return Ok(ResponseEvent::AssistantText(content));
+                },
+            }
+        }
+
         loop {
             match self.next().await {
                 Ok(Some(output)) => match output {
                     ChatResponseStream::AssistantResponseEvent { content } => {
                         self.assistant_text.push_str(&content);
-                        let text = self.buffered_line.take();
-                        self.buffered_line = Some(content);
-                        if let Some(text) = text {
-                            return Ok(ResponseEvent::AssistantText(text));
-                        }
-                    },
-                    ChatResponseStream::CodeReferenceEvent(_) => {
-                        self.buffered_line = None;
+                        return Ok(ResponseEvent::AssistantText(content));
                     },
                     ChatResponseStream::InvalidStateEvent { reason, message } => {
                         error!(%reason, %message, "invalid state event");
@@ -141,11 +131,6 @@ impl ResponseParser {
                     _ => {},
                 },
                 Ok(None) => {
-                    if let Some(text) = self.buffered_line.take() {
-                        self.short_circuit = true;
-                        return Ok(ResponseEvent::AssistantText(text));
-                    }
-
                     let message = Message(ChatMessage::AssistantResponseMessage(AssistantResponseMessage {
                         message_id: self.message_id.take(),
                         content: std::mem::take(&mut self.assistant_text),
@@ -248,6 +233,10 @@ mod tests {
             ChatResponseStream::AssistantResponseEvent {
                 content: " there".to_string(),
             },
+            ChatResponseStream::AssistantResponseEvent {
+                content: "IGNORE ME PLEASE".to_string(),
+            },
+            ChatResponseStream::CodeReferenceEvent(()),
             ChatResponseStream::ToolUseEvent {
                 tool_use_id: tool_use_id.clone(),
                 name: tool_name.clone(),