From 68e0105627aa21679a7579e2c061a2f9bbd6b843 Mon Sep 17 00:00:00 2001
From: Mani Rash Ahmadi <66619933+captaindpt@users.noreply.github.com>
Date: Mon, 28 Apr 2025 11:25:11 -0400
Subject: [PATCH] Agent: Include partial output if terminal tool fails (#29115)

This PR addresses the behavior of the agent's terminal tool when the
executed command is interrupted or fails after producing some output.
Currently, if the command doesn't finish successfully, any partial
output captured before the interruption/failure is discarded, and only
an error message (or a generic cancellation message) is returned to the
LLM.

This change modifies the `run_command_limited` function in the terminal
tool to catch errors when awaiting the command's status (which includes
interruptions). In the case of such an error, it now includes any
partial stdout/stderr captured up to that point within the error message
returned to the `ToolUseState`. This ensures the LLM receives the
partial context even when the command doesn't complete cleanly, framed
appropriately as part of an error/interruption message.

Closes #29101

Release Notes:

- N/A
---
 crates/assistant_tools/src/terminal_tool.rs | 75 ++++++++++++---------
 1 file changed, 44 insertions(+), 31 deletions(-)

diff --git a/crates/assistant_tools/src/terminal_tool.rs b/crates/assistant_tools/src/terminal_tool.rs
index 68f0a64d1477a58b7652dd2629386a4b631fc7d9..bb9d1d93dcac412d095e7cb2c6b78eb3912fc430 100644
--- a/crates/assistant_tools/src/terminal_tool.rs
+++ b/crates/assistant_tools/src/terminal_tool.rs
@@ -202,39 +202,52 @@ async fn run_command_limited(working_dir: Arc<Path>, command: String) -> Result<
     consume_reader(out_reader, truncated).await?;
     consume_reader(err_reader, truncated).await?;
 
-    let status = cmd.status().await.context("Failed to get command status")?;
-
-    let output_string = if truncated {
-        // Valid to find `\n` in UTF-8 since 0-127 ASCII characters are not used in
-        // multi-byte characters.
-        let last_line_ix = combined_buffer.bytes().rposition(|b| b == b'\n');
-        let combined_buffer = &combined_buffer[..last_line_ix.unwrap_or(combined_buffer.len())];
-
-        format!(
-            "Command output too long. The first {} bytes:\n\n{}",
-            combined_buffer.len(),
-            output_block(&combined_buffer),
-        )
-    } else {
-        output_block(&combined_buffer)
-    };
+    // Handle potential errors during status retrieval, including interruption.
+    match cmd.status().await {
+        Ok(status) => {
+            let output_string = if truncated {
+                // Valid to find `\n` in UTF-8 since 0-127 ASCII characters are not used in
+                // multi-byte characters.
+                let last_line_ix = combined_buffer.bytes().rposition(|b| b == b'\n');
+                let buffer_content =
+                    &combined_buffer[..last_line_ix.unwrap_or(combined_buffer.len())];
+
+                format!(
+                    "Command output too long. The first {} bytes:\n\n{}",
+                    buffer_content.len(),
+                    output_block(buffer_content),
+                )
+            } else {
+                output_block(&combined_buffer)
+            };
 
-    let output_with_status = if status.success() {
-        if output_string.is_empty() {
-            "Command executed successfully.".to_string()
-        } else {
-            output_string.to_string()
-        }
-    } else {
-        format!(
-            "Command failed with exit code {} (shell: {}).\n\n{}",
-            status.code().unwrap_or(-1),
-            shell,
-            output_string,
-        )
-    };
+            let output_with_status = if status.success() {
+                if output_string.is_empty() {
+                    "Command executed successfully.".to_string()
+                } else {
+                    output_string
+                }
+            } else {
+                format!(
+                    "Command failed with exit code {} (shell: {}).\n\n{}",
+                    status.code().unwrap_or(-1),
+                    shell,
+                    output_string,
+                )
+            };
 
-    Ok(output_with_status)
+            Ok(output_with_status)
+        }
+        Err(err) => {
+            // Error occurred getting status (potential interruption). Include partial output.
+            let partial_output = output_block(&combined_buffer);
+            let error_message = format!(
+                "Command failed or was interrupted.\nPartial output captured:\n\n{}",
+                partial_output
+            );
+            Err(anyhow!(err).context(error_message))
+        }
+    }
 }
 
 async fn consume_reader<T: AsyncReadExt + Unpin>(