From cc8fe25736e593c391675e61bc108f332da43851 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 d777d5cb0ec00d590195e9a1560ea0ab559948d6..c2863a2b442a809f58c7f4ae1bcad94ab0c7ed37 100644
--- a/crates/assistant_tools/src/terminal_tool.rs
+++ b/crates/assistant_tools/src/terminal_tool.rs
@@ -205,39 +205,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>(