Feat: Support streaming of execute bash commands

chaynabors · chaynabors · commit c306fab6b56c · 2025-03-03T22:11:56.000-08:00
diff --git a/crates/q_cli/src/cli/chat/conversation_state.rs b/crates/q_cli/src/cli/chat/conversation_state.rs
@@ -32,6 +32,7 @@ use tracing::{
 };
 
 use super::tools::ToolSpec;
+use super::truncate_safe;
 use crate::cli::chat::tools::{
     InputSchema,
     InvokeOutput,
@@ -86,7 +87,7 @@ impl ConversationState {
         self.history.clear();
     }
 
-    pub async fn append_new_user_message(&mut self, input: String) {
+    pub fn append_new_user_message(&mut self, input: String) {
         debug_assert!(self.next_message.is_none(), "next_message should not exist");
         if let Some(next_message) = self.next_message.as_ref() {
             warn!(?next_message, "next_message should not exist");
@@ -338,24 +339,6 @@ fn build_shell_state() -> ShellState {
     }
 }
 
-fn truncate_safe(s: &str, max_bytes: usize) -> &str {
-    if s.len() <= max_bytes {
-        return s;
-    }
-
-    let mut byte_count = 0;
-    let mut char_indices = s.char_indices();
-
-    for (byte_idx, _) in &mut char_indices {
-        if byte_count + (byte_idx - byte_count) > max_bytes {
-            break;
-        }
-        byte_count = byte_idx;
-    }
-
-    &s[..byte_count]
-}
-
 #[cfg(test)]
 mod tests {
     use fig_api_client::model::{
@@ -388,7 +371,7 @@ mod tests {
 
         // First, build a large conversation history. We need to ensure that the order is always
         // User -> Assistant -> User -> Assistant ...and so on.
-        conversation_state.append_new_user_message("start".to_string()).await;
+        conversation_state.append_new_user_message("start".to_string());
         for i in 0..=100 {
             let s = conversation_state.as_sendable_conversation_state();
             assert!(
@@ -408,7 +391,7 @@ mod tests {
                 content: i.to_string(),
                 tool_uses: None,
             });
-            conversation_state.append_new_user_message(i.to_string()).await;
+            conversation_state.append_new_user_message(i.to_string());
         }
 
         let s = conversation_state.as_sendable_conversation_state();
@@ -441,7 +424,7 @@ mod tests {
         let mut conversation_state = ConversationState::new(load_tools().unwrap());
 
         // Build a long conversation history of tool use results.
-        conversation_state.append_new_user_message("start".to_string()).await;
+        conversation_state.append_new_user_message("start".to_string());
         for i in 0..=(MAX_CONVERSATION_STATE_HISTORY_LEN + 100) {
             let _ = conversation_state.as_sendable_conversation_state();
             conversation_state.push_assistant_message(AssistantResponseMessage {
diff --git a/crates/q_cli/src/cli/chat/mod.rs b/crates/q_cli/src/cli/chat/mod.rs
@@ -464,7 +464,7 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
                 if !tool_uses.is_empty() {
                     self.conversation_state.abandon_tool_use(tool_uses, user_input);
                 } else {
-                    self.conversation_state.append_new_user_message(user_input).await;
+                    self.conversation_state.append_new_user_message(user_input);
                 }
 
                 self.send_tool_use_telemetry().await;
@@ -874,6 +874,24 @@ Hi, I'm <g>Amazon Q</g>. Ask me anything.
     }
 }
 
+pub fn truncate_safe(s: &str, max_bytes: usize) -> &str {
+    if s.len() <= max_bytes {
+        return s;
+    }
+
+    let mut byte_count = 0;
+    let mut char_indices = s.char_indices();
+
+    for (byte_idx, _) in &mut char_indices {
+        if byte_count + (byte_idx - byte_count) > max_bytes {
+            break;
+        }
+        byte_count = byte_idx;
+    }
+
+    &s[..byte_count]
+}
+
 #[derive(Debug)]
 struct ToolUseEventBuilder {
     pub conversation_id: String,
diff --git a/crates/q_cli/src/cli/chat/tools/execute_bash.rs b/crates/q_cli/src/cli/chat/tools/execute_bash.rs
@@ -1,7 +1,7 @@
+use std::collections::VecDeque;
 use std::io::Write;
 use std::process::Stdio;
 
-use bstr::ByteSlice;
 use crossterm::style::{
     self,
     Color,
@@ -16,77 +16,113 @@ use eyre::{
 };
 use fig_os_shim::Context;
 use serde::Deserialize;
+use tokio::io::AsyncBufReadExt;
+use tokio::select;
+use tracing::error;
 
 use super::{
     InvokeOutput,
     MAX_TOOL_RESPONSE_SIZE,
     OutputKind,
 };
+use crate::cli::chat::truncate_safe;
 
 #[derive(Debug, Clone, Deserialize)]
 pub struct ExecuteBash {
     pub command: String,
-    pub interactive: Option<bool>,
 }
 
 impl ExecuteBash {
     pub async fn invoke(&self, mut updates: impl Write) -> Result<InvokeOutput> {
-        queue!(
+        execute!(
             updates,
             style::SetForegroundColor(Color::Green),
             style::Print(format!("Executing `{}`", &self.command)),
             style::ResetColor,
             style::Print("\n"),
         )?;
 
-        let (stdout, stderr) = match self.interactive {
-            Some(true) => (Stdio::inherit(), Stdio::inherit()),
-            _ => (Stdio::piped(), Stdio::piped()),
-        };
-
-        let output = tokio::process::Command::new("bash")
+        // We need to maintain a handle on stderr and stdout, but pipe it to the terminal as well
+        let mut child = tokio::process::Command::new("bash")
             .arg("-c")
             .arg(&self.command)
             .stdin(Stdio::inherit())
-            .stdout(stdout)
-            .stderr(stderr)
+            .stdout(Stdio::piped())
+            .stderr(Stdio::piped())
             .spawn()
-            .wrap_err_with(|| format!("Unable to spawn command '{}'", &self.command))?
-            .wait_with_output()
-            .await
-            .wrap_err_with(|| format!("Unable to wait on subprocess for command '{}'", &self.command))?;
-        let status = output.status.code().unwrap_or(0).to_string();
-        let stdout = output.stdout.to_str_lossy();
-        let stderr = output.stderr.to_str_lossy();
-
-        if let Some(false) = self.interactive {
-            execute!(updates, style::Print(&stdout))?;
+            .wrap_err_with(|| format!("Unable to spawn command '{}'", &self.command))?;
+
+        let stdout = child.stdout.take().unwrap();
+        let stdout = tokio::io::BufReader::new(stdout);
+        let mut stdout = stdout.lines();
+
+        let stderr = child.stderr.take().unwrap();
+        let stderr = tokio::io::BufReader::new(stderr);
+        let mut stderr = stderr.lines();
+
+        const LINE_COUNT: usize = 1024;
+        let mut stdout_buf = VecDeque::with_capacity(LINE_COUNT);
+        let mut stderr_buf = VecDeque::with_capacity(LINE_COUNT);
+
+        let exit_status = loop {
+            child.stdin.take();
+
+            select! {
+                biased;
+                line = stdout.next_line() => match line {
+                    Ok(Some(line)) => {
+                        writeln!(updates, "{line}")?;
+                        if stdout_buf.len() >= LINE_COUNT {
+                            stdout_buf.pop_front();
+                        }
+                        stdout_buf.push_back(line);
+                    },
+                    Ok(None) => {},
+                    Err(err) => error!(%err, "Failed to read stdout of child process"),
+                },
+                line = stderr.next_line() => match line {
+                    Ok(Some(line)) => {
+                        writeln!(updates, "{line}")?;
+                        if stderr_buf.len() >= LINE_COUNT {
+                            stderr_buf.pop_front();
+                        }
+                        stderr_buf.push_back(line);
+                    },
+                    Ok(None) => {},
+                    Err(err) => error!(%err, "Failed to read stderr of child process"),
+                },
+                exit_status = child.wait() => {
+                    break exit_status;
+                },
+            };
         }
+        .wrap_err_with(|| format!("No exit status for '{}'", &self.command))?;
 
-        let stdout = format!(
-            "{}{}",
-            &stdout[0..stdout.len().min(MAX_TOOL_RESPONSE_SIZE / 3)],
-            if stdout.len() > MAX_TOOL_RESPONSE_SIZE / 3 {
-                " ... truncated"
-            } else {
-                ""
-            }
-        );
-
-        let stderr = format!(
-            "{}{}",
-            &stderr[0..stderr.len().min(MAX_TOOL_RESPONSE_SIZE / 3)],
-            if stderr.len() > MAX_TOOL_RESPONSE_SIZE / 3 {
-                " ... truncated"
-            } else {
-                ""
-            }
-        );
+        updates.flush()?;
+
+        let stdout = stdout_buf.into_iter().collect::<Vec<_>>().join("\n");
+        let stderr = stderr_buf.into_iter().collect::<Vec<_>>().join("\n");
 
         let output = serde_json::json!({
-            "exit_status": status,
-            "stdout": stdout,
-            "stderr": stderr,
+            "exit_status": exit_status.code().unwrap_or(0).to_string(),
+            "stdout": format!(
+                "{}{}",
+                truncate_safe(&stdout, MAX_TOOL_RESPONSE_SIZE / 3),
+                if stdout.len() > MAX_TOOL_RESPONSE_SIZE / 3 {
+                    " ... truncated"
+                } else {
+                    ""
+                }
+            ),
+            "stderr": format!(
+                "{}{}",
+                truncate_safe(&stderr, MAX_TOOL_RESPONSE_SIZE / 3),
+                if stderr.len() > MAX_TOOL_RESPONSE_SIZE / 3 {
+                    " ... truncated"
+                } else {
+                    ""
+                }
+            ),
         });
 
         Ok(InvokeOutput {
@@ -127,7 +163,6 @@ mod tests {
         // Verifying stdout
         let v = serde_json::json!({
             "command": "echo Hello, world!",
-            "interactive": false
         });
         let out = serde_json::from_value::<ExecuteBash>(v)
             .unwrap()
@@ -137,7 +172,7 @@ mod tests {
 
         if let OutputKind::Json(json) = out.output {
             assert_eq!(json.get("exit_status").unwrap(), &0.to_string());
-            assert_eq!(json.get("stdout").unwrap(), "Hello, world!\n");
+            assert_eq!(json.get("stdout").unwrap(), "Hello, world!");
             assert_eq!(json.get("stderr").unwrap(), "");
         } else {
             panic!("Expected JSON output");
@@ -157,7 +192,7 @@ mod tests {
         if let OutputKind::Json(json) = out.output {
             assert_eq!(json.get("exit_status").unwrap(), &0.to_string());
             assert_eq!(json.get("stdout").unwrap(), "");
-            assert_eq!(json.get("stderr").unwrap(), "Hello, world!\n");
+            assert_eq!(json.get("stderr").unwrap(), "Hello, world!");
         } else {
             panic!("Expected JSON output");
         }
diff --git a/crates/q_cli/src/cli/chat/tools/tool_index.json b/crates/q_cli/src/cli/chat/tools/tool_index.json
@@ -8,13 +8,9 @@
         "command": {
           "type": "string",
           "description": "Bash command to execute"
-        },
-        "interactive": {
-          "type": "boolean",
-          "description": "Whether or not the command is interactive. Interactive commands like nano will overtake our conversation until exited. On exit, they will have produced no stderr or stdout."
         }
       },
-      "required": ["command", "interactive"]
+      "required": ["command"]
     }
   },
   "fs_read": {
