Fixed: Preserve buffered output on tokio bash timeout

Sewer56 · Sewer56 · commit dd56e3b57869 · 2026-02-27T23:02:46.000Z
Avoid canceling stdout/stderr drains by spawning independent pipe-read tasks and racing only child completion against timeout. On timeout, kill the process tree, await drain tasks briefly to capture buffered output, and include captured stdout/stderr in timeout errors; add a regression test for pre-timeout output.
diff --git a/src/llm-coding-tools-core/src/tools/bash/tokio_impl.rs b/src/llm-coding-tools-core/src/tools/bash/tokio_impl.rs
@@ -2,11 +2,78 @@
 
 use super::{BashOutput, PIPE_BUFFER_CAPACITY};
 use crate::error::{ToolError, ToolResult};
+use core::fmt::Write;
 use process_wrap::tokio::*;
 use std::path::Path;
+use std::pin::Pin;
 use std::process::Stdio;
 use std::time::Duration;
-use tokio::io::AsyncReadExt;
+use tokio::io::{AsyncRead, AsyncReadExt};
+use tokio::task::JoinHandle;
+
+/// Maximum time to wait for pipe drains after timeout kill.
+const PIPE_DRAIN_GRACE_PERIOD: Duration = Duration::from_millis(100);
+
+#[inline]
+fn spawn_pipe_drain_task<R>(mut pipe: R) -> JoinHandle<Vec<u8>>
+where
+    R: AsyncRead + Unpin + Send + 'static,
+{
+    tokio::spawn(async move {
+        let mut buf = Vec::with_capacity(PIPE_BUFFER_CAPACITY);
+        let _ = pipe.read_to_end(&mut buf).await;
+        buf
+    })
+}
+
+#[inline]
+async fn await_pipe_drain_task(task: JoinHandle<Vec<u8>>) -> Vec<u8> {
+    task.await.unwrap_or_default()
+}
+
+#[inline]
+async fn await_pipe_drain_task_with_grace(
+    mut task: JoinHandle<Vec<u8>>,
+    grace: Duration,
+) -> Vec<u8> {
+    tokio::select! {
+        result = &mut task => result.unwrap_or_default(),
+        _ = tokio::time::sleep(grace) => {
+            task.abort();
+            let _ = task.await;
+            Vec::new()
+        }
+    }
+}
+
+#[inline]
+fn timeout_with_buffered_output(
+    timeout: Duration,
+    stdout_data: &[u8],
+    stderr_data: &[u8],
+) -> ToolError {
+    let stdout = String::from_utf8_lossy(stdout_data);
+    let stderr = String::from_utf8_lossy(stderr_data);
+
+    // Base message + outputs + stderr label.
+    let mut message = String::with_capacity(stdout.len() + stderr.len() + 64);
+    let _ = write!(message, "command timed out after {}ms", timeout.as_millis());
+
+    if !stdout.is_empty() {
+        message.push('\n');
+        message.push_str(&stdout);
+    }
+
+    if !stderr.is_empty() {
+        if stdout.is_empty() || !stdout.ends_with('\n') {
+            message.push('\n');
+        }
+        message.push_str("[stderr]\n");
+        message.push_str(&stderr);
+    }
+
+    ToolError::Timeout(message)
+}
 
 /// Executes a shell command with optional working directory and timeout.
 ///
@@ -67,40 +134,29 @@ pub async fn execute_command(
 
     // Take stdout/stderr handles to drain them concurrently with process wait.
     // This prevents deadlock when output exceeds pipe buffer (~64KB Linux, ~4KB Windows).
-    let mut stdout_pipe = child.stdout().take().expect("stdout was piped");
-    let mut stderr_pipe = child.stderr().take().expect("stderr was piped");
+    let stdout_pipe = child.stdout().take().expect("stdout was piped");
+    let stderr_pipe = child.stderr().take().expect("stderr was piped");
 
-    // Race between timeout and (process completion + pipe draining).
-    // Using join! inside select! avoids tokio::spawn overhead while still
-    // providing concurrent I/O for the pipe reads.
-    tokio::select! {
+    // Keep output drains independent from timeout selection so timed-out
+    // commands can still return buffered stdout/stderr.
+    let stdout_task = spawn_pipe_drain_task(stdout_pipe);
+    let stderr_task = spawn_pipe_drain_task(stderr_pipe);
+
+    // Race between timeout and process completion. Pipe drain tasks keep running
+    // regardless of which branch wins this select.
+    let wait_result = tokio::select! {
         biased;  // Check timeout first for consistent behavior
 
-        _ = tokio::time::sleep(timeout) => {
-            // Timeout: explicitly kill the process tree (Job Object on Windows, process group on Unix)
-            let _ = Pin::from(child.kill()).await;
-            Err(ToolError::Timeout(format!(
-                "command timed out after {}ms",
-                timeout.as_millis()
-            )))
-        }
+        _ = tokio::time::sleep(timeout) => None,
+        status = child.wait() => Some(status),
+    };
 
-        result = async {
-            tokio::join!(
-                child.wait(),
-                async {
-                    let mut buf = Vec::with_capacity(PIPE_BUFFER_CAPACITY);
-                    let _ = stdout_pipe.read_to_end(&mut buf).await;
-                    buf
-                },
-                async {
-                    let mut buf = Vec::with_capacity(PIPE_BUFFER_CAPACITY);
-                    let _ = stderr_pipe.read_to_end(&mut buf).await;
-                    buf
-                }
-            )
-        } => {
-            let (status, stdout_data, stderr_data) = result;
+    match wait_result {
+        Some(status) => {
+            let (stdout_data, stderr_data) = tokio::join!(
+                await_pipe_drain_task(stdout_task),
+                await_pipe_drain_task(stderr_task)
+            );
             let status = status.map_err(|e| ToolError::Execution(e.to_string()))?;
 
             Ok(BashOutput {
@@ -109,6 +165,22 @@ pub async fn execute_command(
                 stderr: String::from_utf8_lossy(&stderr_data).into_owned(),
             })
         }
+        None => {
+            // Timeout: explicitly kill the process tree (Job Object on Windows,
+            // process group on Unix), then briefly await pipe drains for buffered output.
+            let _ = Pin::from(child.kill()).await;
+
+            let (stdout_data, stderr_data) = tokio::join!(
+                await_pipe_drain_task_with_grace(stdout_task, PIPE_DRAIN_GRACE_PERIOD),
+                await_pipe_drain_task_with_grace(stderr_task, PIPE_DRAIN_GRACE_PERIOD)
+            );
+
+            Err(timeout_with_buffered_output(
+                timeout,
+                &stdout_data,
+                &stderr_data,
+            ))
+        }
     }
 }
 
@@ -157,6 +229,24 @@ mod tests {
         assert!(matches!(result, Err(ToolError::Timeout(_))));
     }
 
+    #[tokio::test]
+    async fn timeout_preserves_buffered_output() {
+        let cmd = if cfg!(target_os = "windows") {
+            "echo stdout-before-timeout & echo stderr-before-timeout 1>&2 & ping -n 10 127.0.0.1 >nul"
+        } else {
+            "echo stdout-before-timeout; echo stderr-before-timeout 1>&2; sleep 10"
+        };
+
+        let result = execute_command(cmd, None, Duration::from_millis(100)).await;
+        match result {
+            Err(ToolError::Timeout(message)) => {
+                assert!(message.contains("stdout-before-timeout"));
+                assert!(message.contains("stderr-before-timeout"));
+            }
+            other => panic!("expected timeout error, got: {other:?}"),
+        }
+    }
+
     #[tokio::test]
     async fn invalid_workdir_returns_error() {
         let result = execute_command(
