From 86527e19a5a78213b9044283a51549fce3cdaac6 Mon Sep 17 00:00:00 2001 From: Marks Date: Sun, 8 Feb 2026 03:25:08 +0000 Subject: [PATCH 1/2] Enforce shell tool timeout and kill process group on expiry The shell tool parsed the timeout parameter but never used it -- commands ran via std::process::Command::output() which blocks indefinitely. This caused the agent to freeze whenever the LLM launched a long-running or persistent process (e.g. a daemon). Switch to tokio::process::Command with tokio::time::timeout() for actual enforcement. Spawn children in their own process group (.process_group(0)) so kill(-pgid, SIGKILL) reaches all descendants on timeout. Update the tool description to warn the LLM against launching background daemons. Fixes AnthonyRonning/sage#8 --- Cargo.lock | 1 + crates/sage-core/Cargo.toml | 1 + crates/sage-core/src/shell_tool.rs | 70 +++++++++++++++++++++++++----- 3 files changed, 61 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1205955..8846260 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3350,6 +3350,7 @@ dependencies = [ "diesel_migrations", "dotenvy", "dspy-rs", + "libc", "pgvector", "redis", "reqwest", diff --git a/crates/sage-core/Cargo.toml b/crates/sage-core/Cargo.toml index 9225bd5..f75c802 100644 --- a/crates/sage-core/Cargo.toml +++ b/crates/sage-core/Cargo.toml @@ -45,6 +45,7 @@ cron.workspace = true uuid.workspace = true dotenvy.workspace = true socket2 = "0.5" +libc = "0.2" [dev-dependencies] tokio-test.workspace = true diff --git a/crates/sage-core/src/shell_tool.rs b/crates/sage-core/src/shell_tool.rs index 29e90f3..f772175 100644 --- a/crates/sage-core/src/shell_tool.rs +++ b/crates/sage-core/src/shell_tool.rs @@ -1,11 +1,14 @@ //! Shell command execution tool //! //! Allows Sage to execute arbitrary shell commands within its container. +//! Commands are run asynchronously with enforced timeouts. On timeout the +//! entire process group is killed so that child/background processes cannot +//! outlive the tool invocation and block the agent loop. use anyhow::Result; use async_trait::async_trait; use std::collections::HashMap; -use std::process::Command; +use tokio::process::Command; use tracing::{debug, info, warn}; use crate::sage_agent::{Tool, ToolResult}; @@ -83,7 +86,7 @@ impl Tool for ShellTool { } fn description(&self) -> &str { - "Execute a shell command in the workspace. Has access to CLI tools: git, curl, jq, grep, sed, awk, python3, node, etc. Use for file operations, running scripts, or system commands." + "Execute a shell command in the workspace. Has access to CLI tools: git, curl, jq, grep, sed, awk, python3, node, etc. Use for file operations, running scripts, or system commands. Commands MUST complete within the timeout -- do NOT launch persistent/background daemons (they will be killed when the timeout expires)." } fn args_schema(&self) -> &str { @@ -95,7 +98,7 @@ impl Tool for ShellTool { .get("command") .ok_or_else(|| anyhow::anyhow!("'command' argument is required"))?; - let timeout: u64 = args + let timeout_secs: u64 = args .get("timeout") .and_then(|v| v.parse().ok()) .unwrap_or(DEFAULT_TIMEOUT) @@ -103,7 +106,7 @@ impl Tool for ShellTool { info!( "Executing shell command: {} (timeout: {}s)", - command, timeout + command, timeout_secs ); // Check for blocked patterns @@ -119,16 +122,35 @@ impl Tool for ShellTool { // Ensure workspace exists std::fs::create_dir_all(&self.workspace).ok(); - // Execute command via bash - let result = Command::new("bash") + // Spawn command in a new process group so we can kill the entire tree + // (including any child/background processes) on timeout. + let child = match Command::new("bash") .args(["-c", command]) .current_dir(&self.workspace) .env("HOME", &self.workspace) .env("PWD", &self.workspace) - .output(); + .process_group(0) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .spawn() + { + Ok(child) => child, + Err(e) => { + return Ok(ToolResult { + success: false, + output: String::new(), + error: Some(format!("Failed to execute command: {}", e)), + }); + } + }; + + let timeout_duration = std::time::Duration::from_secs(timeout_secs); - match result { - Ok(output) => { + // Grab the PID now -- wait_with_output() consumes the child. + let child_pid = child.id(); + + match tokio::time::timeout(timeout_duration, child.wait_with_output()).await { + Ok(Ok(output)) => { let stdout = String::from_utf8_lossy(&output.stdout); let stderr = String::from_utf8_lossy(&output.stderr); let exit_code = output.status.code().unwrap_or(-1); @@ -159,11 +181,37 @@ impl Tool for ShellTool { }, }) } - Err(e) => Ok(ToolResult { + Ok(Err(e)) => Ok(ToolResult { success: false, output: String::new(), - error: Some(format!("Failed to execute command: {}", e)), + error: Some(format!("Failed to read command output: {}", e)), }), + Err(_) => { + // Timeout -- kill the entire process group + warn!( + "Shell command timed out after {}s, killing process group: {}", + timeout_secs, command + ); + + if let Some(pid) = child_pid { + let pgid = pid as i32; + // SIGKILL the entire process group (negative pid) + unsafe { + libc::kill(-pgid, libc::SIGKILL); + } + } + + Ok(ToolResult { + success: false, + output: format!( + "Command timed out after {} seconds and was killed (including all child processes). \ + Do NOT launch persistent daemons or background services via this tool -- \ + they will always be killed at timeout.", + timeout_secs + ), + error: Some(format!("Command timed out after {}s", timeout_secs)), + }) + } } } } From 0c7b5b794d448af179d1b16874764027c0fbe792 Mon Sep 17 00:00:00 2001 From: Marks Date: Sun, 8 Feb 2026 21:09:50 +0000 Subject: [PATCH 2/2] Return partial output on timeout, let agent manage its own timeouts Address review feedback: - On timeout, drain stdout/stderr pipes before returning so the agent sees what the command produced before it was killed - Remove 300s hard cap on timeout (now 24h safety rail); agent sets whatever timeout is appropriate, default remains 60s - Clean up tool description: no more lecturing about daemons - Reap zombie process after SIGKILL to avoid leaking PIDs - Extract drain_pipe/drain_stderr/format_output helpers --- crates/sage-core/src/sage_agent.rs | 4 +- crates/sage-core/src/shell_tool.rs | 129 +++++++++++++++++++++-------- 2 files changed, 97 insertions(+), 36 deletions(-) diff --git a/crates/sage-core/src/sage_agent.rs b/crates/sage-core/src/sage_agent.rs index 35c111d..041b2ab 100644 --- a/crates/sage-core/src/sage_agent.rs +++ b/crates/sage-core/src/sage_agent.rs @@ -423,8 +423,8 @@ impl ToolRegistry { // -- Shell tool -- registry.register_descriptor( "shell", - "Execute a shell command in the workspace. Has access to CLI tools: git, curl, jq, grep, sed, awk, python3, node, etc. Use for file operations, running scripts, or system commands.", - r#"{"command": "shell command to execute (supports pipes, redirects)", "timeout": "optional timeout in seconds (default 60, max 300)"}"#, + "Execute a shell command in the workspace. Has access to CLI tools: git, curl, jq, grep, sed, awk, python3, node, etc. Use for file operations, running scripts, or system commands. Set the timeout parameter appropriately for each command (default 60s). If the command exceeds the timeout it will be killed and any partial output returned.", + r#"{"command": "shell command to execute (supports pipes, redirects)", "timeout": "optional timeout in seconds (default 60, set appropriately for long-running commands)"}"#, ); // -- Web search tool -- diff --git a/crates/sage-core/src/shell_tool.rs b/crates/sage-core/src/shell_tool.rs index f772175..d91176e 100644 --- a/crates/sage-core/src/shell_tool.rs +++ b/crates/sage-core/src/shell_tool.rs @@ -4,10 +4,14 @@ //! Commands are run asynchronously with enforced timeouts. On timeout the //! entire process group is killed so that child/background processes cannot //! outlive the tool invocation and block the agent loop. +//! +//! When a command is killed due to timeout, any partial stdout/stderr captured +//! before the kill is included in the result so the agent can see what happened. use anyhow::Result; use async_trait::async_trait; use std::collections::HashMap; +use tokio::io::AsyncReadExt; use tokio::process::Command; use tracing::{debug, info, warn}; @@ -35,8 +39,8 @@ const MAX_OUTPUT_SIZE: usize = 100_000; // 100KB /// Default timeout in seconds const DEFAULT_TIMEOUT: u64 = 60; -/// Maximum timeout in seconds -const MAX_TIMEOUT: u64 = 300; +/// Maximum timeout in seconds (safety rail for clearly nonsensical values) +const MAX_TIMEOUT: u64 = 86_400; // 24 hours /// Shell command execution tool pub struct ShellTool { @@ -59,6 +63,49 @@ impl ShellTool { .copied() } + /// Read all available bytes from an optional pipe handle. + /// Returns the content as a String (lossy UTF-8). + async fn drain_pipe(pipe: &mut Option) -> String { + // This generic approach won't work for ChildStderr directly, so we + // have a separate overload below. Rust doesn't support trait-object + // generics ergonomically here, so we just duplicate for the two types. + if let Some(ref mut handle) = pipe { + let mut buf = Vec::new(); + let _ = handle.read_to_end(&mut buf).await; + String::from_utf8_lossy(&buf).into_owned() + } else { + String::new() + } + } + + /// Read all available bytes from an optional stderr pipe handle. + async fn drain_stderr(pipe: &mut Option) -> String { + if let Some(ref mut handle) = pipe { + let mut buf = Vec::new(); + let _ = handle.read_to_end(&mut buf).await; + String::from_utf8_lossy(&buf).into_owned() + } else { + String::new() + } + } + + /// Build the standard output string from stdout, stderr, and exit code. + fn format_output(&self, stdout: &str, stderr: &str, exit_code: i32) -> String { + let mut result_parts = Vec::new(); + + if !stdout.is_empty() { + result_parts.push(format!("STDOUT:\n{}", stdout.trim())); + } + + if !stderr.is_empty() { + result_parts.push(format!("STDERR:\n{}", stderr.trim())); + } + + result_parts.push(format!("EXIT CODE: {}", exit_code)); + + self.truncate_output(result_parts.join("\n\n")) + } + /// Truncate output if too long (handles UTF-8 boundaries safely) fn truncate_output(&self, output: String) -> String { if output.len() > MAX_OUTPUT_SIZE { @@ -86,11 +133,11 @@ impl Tool for ShellTool { } fn description(&self) -> &str { - "Execute a shell command in the workspace. Has access to CLI tools: git, curl, jq, grep, sed, awk, python3, node, etc. Use for file operations, running scripts, or system commands. Commands MUST complete within the timeout -- do NOT launch persistent/background daemons (they will be killed when the timeout expires)." + "Execute a shell command in the workspace. Has access to CLI tools: git, curl, jq, grep, sed, awk, python3, node, etc. Use for file operations, running scripts, or system commands. Set the timeout parameter appropriately for each command (default 60s). If the command exceeds the timeout it will be killed and any partial output returned." } fn args_schema(&self) -> &str { - r#"{"command": "shell command to execute (supports pipes, redirects)", "timeout": "optional timeout in seconds (default 60, max 300)"}"# + r#"{"command": "shell command to execute (supports pipes, redirects)", "timeout": "optional timeout in seconds (default 60, set appropriately for long-running commands)"}"# } async fn execute(&self, args: &HashMap) -> Result { @@ -124,7 +171,7 @@ impl Tool for ShellTool { // Spawn command in a new process group so we can kill the entire tree // (including any child/background processes) on timeout. - let child = match Command::new("bash") + let mut child = match Command::new("bash") .args(["-c", command]) .current_dir(&self.workspace) .env("HOME", &self.workspace) @@ -146,35 +193,30 @@ impl Tool for ShellTool { let timeout_duration = std::time::Duration::from_secs(timeout_secs); - // Grab the PID now -- wait_with_output() consumes the child. + // Take ownership of the pipe handles so we can read partial output on + // timeout. child.wait() only waits for exit -- it does not consume the + // pipes, unlike child.wait_with_output(). + let mut child_stdout = child.stdout.take(); + let mut child_stderr = child.stderr.take(); + // Note: child_stdout is Option, child_stderr is Option. + // We use separate drain helpers because they are different types. let child_pid = child.id(); - match tokio::time::timeout(timeout_duration, child.wait_with_output()).await { - Ok(Ok(output)) => { - let stdout = String::from_utf8_lossy(&output.stdout); - let stderr = String::from_utf8_lossy(&output.stderr); - let exit_code = output.status.code().unwrap_or(-1); - - let mut result_parts = Vec::new(); - - if !stdout.is_empty() { - result_parts.push(format!("STDOUT:\n{}", stdout.trim())); - } - - if !stderr.is_empty() { - result_parts.push(format!("STDERR:\n{}", stderr.trim())); - } + match tokio::time::timeout(timeout_duration, child.wait()).await { + Ok(Ok(status)) => { + // Command finished within the timeout -- drain remaining output. + let stdout = Self::drain_pipe(&mut child_stdout).await; + let stderr = Self::drain_stderr(&mut child_stderr).await; + let exit_code = status.code().unwrap_or(-1); - result_parts.push(format!("EXIT CODE: {}", exit_code)); - - let output_str = self.truncate_output(result_parts.join("\n\n")); + let output_str = self.format_output(&stdout, &stderr, exit_code); debug!("Shell command completed with exit code {}", exit_code); Ok(ToolResult { - success: output.status.success(), + success: status.success(), output: output_str, - error: if output.status.success() { + error: if status.success() { None } else { Some(format!("Command exited with code {}", exit_code)) @@ -184,10 +226,11 @@ impl Tool for ShellTool { Ok(Err(e)) => Ok(ToolResult { success: false, output: String::new(), - error: Some(format!("Failed to read command output: {}", e)), + error: Some(format!("Failed to wait on command: {}", e)), }), Err(_) => { - // Timeout -- kill the entire process group + // Timeout -- kill the entire process group first, then drain + // whatever partial output was written before the kill. warn!( "Shell command timed out after {}s, killing process group: {}", timeout_secs, command @@ -201,14 +244,32 @@ impl Tool for ShellTool { } } + // Reap the zombie so we don't leak it. + let _ = child.wait().await; + + // Drain whatever was buffered in the pipes before the kill. + let stdout = Self::drain_pipe(&mut child_stdout).await; + let stderr = Self::drain_stderr(&mut child_stderr).await; + + let mut result_parts = Vec::new(); + + if !stdout.is_empty() { + result_parts.push(format!("STDOUT (partial):\n{}", stdout.trim())); + } + if !stderr.is_empty() { + result_parts.push(format!("STDERR (partial):\n{}", stderr.trim())); + } + + result_parts.push(format!( + "[Command timed out after {}s and was killed]", + timeout_secs + )); + + let output_str = self.truncate_output(result_parts.join("\n\n")); + Ok(ToolResult { success: false, - output: format!( - "Command timed out after {} seconds and was killed (including all child processes). \ - Do NOT launch persistent daemons or background services via this tool -- \ - they will always be killed at timeout.", - timeout_secs - ), + output: output_str, error: Some(format!("Command timed out after {}s", timeout_secs)), }) }