Add background command execution to run_command

index bccf236..4f302c6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -250,8 +250,8 @@ fn tool_kind_for(tool_name: &str) -> ToolKind { ToolKind::Edit } "delete_file" => ToolKind::Delete, - "run_command" => ToolKind::Execute, - "read_file" | "list_directory" => ToolKind::Read, + "run_command" | "kill_command" => ToolKind::Execute, + "read_file" | "list_directory" | "command_output" => ToolKind::Read, "search_files" | "glob" => ToolKind::Search, "read_website" => ToolKind::Fetch, "write_todos" => ToolKind::Think,

index cf1e6c3..a5d383e 100644 --- a/src/permissions.rs +++ b/src/permissions.rs @@ -67,7 +67,7 @@ pub enum Decision { pub fn classify(tool_name: &str) -> ToolRisk { match tool_name { "read_file" | "list_directory" | "search_files" | "glob" | "read_website" - | "write_todos" | "skill" => ToolRisk::ReadOnly, + | "write_todos" | "skill" | "command_output" => ToolRisk::ReadOnly, _ => ToolRisk::Mutating, } } @@ -247,6 +247,7 @@ mod tests { "read_website", "write_todos", "skill", + "command_output", ] { assert_eq!(classify(tool), ToolRisk::ReadOnly, "{tool}"); assert_eq!(decision_for("t-ro", tool), Decision::Allow, "{tool}"); @@ -262,6 +263,7 @@ mod tests { "create_directory", "delete_file", "run_command", + "kill_command", "remember", "mcp__server__anything", "totally_unknown_tool",

index 721f965..5002a88 100644 --- a/src/tools.rs +++ b/src/tools.rs @@ -2,9 +2,12 @@ use regex::Regex; use serde_json::{Value, json}; +use std::collections::HashMap; use std::fs; use std::path::{Path, PathBuf}; use std::process::Command; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; const WEBSITE_READ_CHAR_LIMIT: usize = 20_000; const WEBSITE_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20); @@ -231,8 +234,10 @@ pub fn all_tools() -> Vec<AgentTool> { use this for `git clone`, `git init`, and normal repo setup steps. \ In smbCloud repos, prefer existing workspace commands, Rails \ conventions, and deploy flows over inventing new command sequences. \ - Commands that run indefinitely (servers, watchers) will be killed \ - after 120 seconds.", + Foreground commands are killed after 120 seconds — run servers, \ + watchers, builds, test suites, and anything that may exceed two \ + minutes with run_in_background set to true, then poll with \ + command_output.", parameters_schema: json!({ "type": "object", "properties": { @@ -243,6 +248,10 @@ pub fn all_tools() -> Vec<AgentTool> { "cwd": { "type": "string", "description": "Working directory for the command. Defaults to \".\" (current directory)." + }, + "run_in_background": { + "type": "boolean", + "description": "Start the command as a background task and return a task id immediately instead of waiting for it to finish. Set this to true for servers, watchers, builds, test suites, and anything that may run longer than two minutes (foreground commands are killed after 120 seconds). Poll the task's output and status with command_output, and stop it with kill_command. Background tasks are killed when sigit exits. Defaults to false." } }, "required": ["command"], @@ -373,6 +382,45 @@ pub fn all_tools() -> Vec<AgentTool> { "additionalProperties": false }), }, + AgentTool { + name: "command_output", + description: "Get the output a background task (started with run_command's \ + run_in_background) has produced since your last check, plus its \ + status: still running, or exited with an exit code. Poll this \ + periodically to follow builds, test suites, and servers. Between \ + polls output is buffered up to 50 000 bytes per task; older \ + output beyond that is dropped and the truncation is noted.", + parameters_schema: json!({ + "type": "object", + "properties": { + "task_id": { + "type": "integer", + "description": "Id of the background task, as returned by run_command with run_in_background." + } + }, + "required": ["task_id"], + "additionalProperties": false + }), + }, + AgentTool { + name: "kill_command", + description: "Kill a background task started with run_command's run_in_background. \ + Reports the tail of the task's unread output and confirms it was \ + killed. Use this to stop servers or watchers you no longer need and \ + runaway commands. Background tasks are also killed automatically \ + when sigit exits.", + parameters_schema: json!({ + "type": "object", + "properties": { + "task_id": { + "type": "integer", + "description": "Id of the background task, as returned by run_command with run_in_background." + } + }, + "required": ["task_id"], + "additionalProperties": false + }), + }, ] } @@ -399,6 +447,8 @@ pub async fn execute_tool(name: &str, arguments: &str) -> String { "remember" => exec_remember(arguments), "delete_file" => exec_delete_file(arguments), "run_command" => exec_run_command(arguments), + "command_output" => exec_command_output(arguments), + "kill_command" => exec_kill_command(arguments), "skill" => crate::skills::activate_skill(arguments), // Tools discovered from MCP servers are namespaced `mcp__<server>__<tool>` // and forwarded to the owning server. @@ -1379,7 +1429,27 @@ fn exec_delete_file(arguments: &str) -> String { const COMMAND_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120); const COMMAND_OUTPUT_LIMIT: usize = 50_000; -/// runs via `sh -c` / `cmd /C`; killed after COMMAND_TIMEOUT. +/// Spawn `command_str` through the platform shell with piped stdout/stderr. +/// Shared by the foreground and background paths of `run_command`. +fn spawn_shell(command_str: &str, cwd_path: &Path) -> std::io::Result<std::process::Child> { + #[cfg(unix)] + let (shell, flag) = ("sh", "-c"); + #[cfg(windows)] + let (shell, flag) = ("cmd", "/C"); + + Command::new(shell) + .arg(flag) + .arg(command_str) + .current_dir(cwd_path) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .spawn() +} + +/// runs via `sh -c` / `cmd /C`; killed after COMMAND_TIMEOUT unless +/// `run_in_background` is set, in which case the child is registered as a +/// background task and polled with `command_output` / stopped with +/// `kill_command`. fn exec_run_command(arguments: &str) -> String { let args: Value = match serde_json::from_str(arguments) { Ok(v) => v, @@ -1403,30 +1473,18 @@ fn exec_run_command(arguments: &str) -> String { return format!("Error: working directory does not exist: {cwd_str}"); } - log::info!("run_command: `{command_str}` in `{cwd_str}`"); + let run_in_background = args + .get("run_in_background") + .and_then(Value::as_bool) + .unwrap_or(false); - #[cfg(unix)] - let mut child = match Command::new("sh") - .arg("-c") - .arg(command_str) - .current_dir(&cwd_path) - .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::piped()) - .spawn() - { - Ok(c) => c, - Err(err) => return format!("Error: failed to spawn command: {err}"), - }; + log::info!("run_command: `{command_str}` in `{cwd_str}` (background: {run_in_background})"); - #[cfg(windows)] - let mut child = match Command::new("cmd") - .arg("/C") - .arg(command_str) - .current_dir(&cwd_path) - .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::piped()) - .spawn() - { + if run_in_background { + return start_background_task(command_str, &cwd_path); + } + + let mut child = match spawn_shell(command_str, &cwd_path) { Ok(c) => c, Err(err) => return format!("Error: failed to spawn command: {err}"), }; @@ -1477,6 +1535,252 @@ fn exec_run_command(arguments: &str) -> String { } } +// ── background tasks (command_output / kill_command) ───────────────────────── + +/// A command started with `run_in_background`. The child is never detached, +/// so background tasks die with the sigit process. +struct BackgroundTask { + command: String, + child: std::process::Child, + /// Combined stdout+stderr drained by the reader threads, shared with them. + output: Arc<Mutex<TaskOutput>>, + /// Exit code cached once observed; `-1` when there is no code (killed by + /// a signal). `None` while the task is still running. + exit_code: Option<i32>, + /// Set when the task was stopped via `kill_command`. + killed: bool, +} + +/// Output accumulated for a background task since the last poll. +#[derive(Default)] +struct TaskOutput { + buf: String, + /// Whether the oldest output was dropped because `buf` hit the cap. + dropped: bool, +} + +/// Process-global background task table (same pattern as `mcp::MCP`). +fn tasks() -> &'static Mutex<HashMap<u64, BackgroundTask>> { + static TASKS: OnceLock<Mutex<HashMap<u64, BackgroundTask>>> = OnceLock::new(); + TASKS.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn lock_tasks() -> std::sync::MutexGuard<'static, HashMap<u64, BackgroundTask>> { + tasks() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +fn next_task_id() -> u64 { + static NEXT_ID: AtomicU64 = AtomicU64::new(1); + NEXT_ID.fetch_add(1, Ordering::Relaxed) +} + +/// Drain a child's stream into the task's shared buffer on a plain std thread, +/// capping the buffer at COMMAND_OUTPUT_LIMIT by dropping the oldest output. +fn spawn_output_reader<R: std::io::Read + Send + 'static>( + mut stream: R, + output: Arc<Mutex<TaskOutput>>, +) { + std::thread::spawn(move || { + let mut chunk = [0u8; 8192]; + loop { + match stream.read(&mut chunk) { + Ok(0) | Err(_) => break, + Ok(n) => { + let text = String::from_utf8_lossy(&chunk[..n]).into_owned(); + let mut out = output + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + out.buf.push_str(&text); + if out.buf.len() > COMMAND_OUTPUT_LIMIT { + let mut cut = out.buf.len() - COMMAND_OUTPUT_LIMIT; + while !out.buf.is_char_boundary(cut) { + cut += 1; + } + out.buf.drain(..cut); + out.dropped = true; + } + } + } + } + }); +} + +/// Background branch of `run_command`: spawn, register, return immediately. +fn start_background_task(command_str: &str, cwd_path: &Path) -> String { + let mut child = match spawn_shell(command_str, cwd_path) { + Ok(c) => c, + Err(err) => return format!("Error: failed to spawn command: {err}"), + }; + + let output = Arc::new(Mutex::new(TaskOutput::default())); + if let Some(stdout) = child.stdout.take() { + spawn_output_reader(stdout, Arc::clone(&output)); + } + if let Some(stderr) = child.stderr.take() { + spawn_output_reader(stderr, Arc::clone(&output)); + } + + let task_id = next_task_id(); + lock_tasks().insert( + task_id, + BackgroundTask { + command: command_str.to_string(), + child, + output, + exit_code: None, + killed: false, + }, + ); + + format!( + "Started background task {task_id}: `{command_str}`. Poll its output and status \ + with command_output (task_id: {task_id}); stop it with kill_command. The task is \ + killed when sigit exits." + ) +} + +/// Check (and cache) whether a task's child has exited. Returns the exit code, +/// or `None` while it is still running. +fn poll_exit_code(task: &mut BackgroundTask) -> Option<i32> { + if task.exit_code.is_none() + && let Ok(Some(status)) = task.child.try_wait() + { + task.exit_code = Some(status.code().unwrap_or(-1)); + } + task.exit_code +} + +/// Take everything the task has printed since the last poll, plus whether +/// older output was dropped at the buffer cap. +fn drain_task_output(task: &BackgroundTask) -> (String, bool) { + let mut out = task + .output + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + ( + std::mem::take(&mut out.buf), + std::mem::take(&mut out.dropped), + ) +} + +fn parse_task_id(arguments: &str) -> Result<u64, String> { + let args: Value = serde_json::from_str(arguments) + .map_err(|err| format!("Error: failed to parse arguments: {err}"))?; + args.get("task_id") + .and_then(Value::as_u64) + .ok_or_else(|| "Error: missing required parameter \"task_id\"".to_string()) +} + +fn unknown_task(task_id: u64) -> String { + format!( + "Error: no background task with id {task_id}. Start one with run_command and \ + run_in_background set to true." + ) +} + +fn dropped_note(dropped: bool) -> &'static str { + if dropped { + "\n(note: earlier output was dropped after exceeding the 50000-byte buffer)" + } else { + "" + } +} + +/// `command_output` tool: output since the last poll + running/exited status. +fn exec_command_output(arguments: &str) -> String { + let task_id = match parse_task_id(arguments) { + Ok(id) => id, + Err(err) => return err, + }; + + let mut map = lock_tasks(); + let Some(task) = map.get_mut(&task_id) else { + return unknown_task(task_id); + }; + + let was_running = task.exit_code.is_none(); + let exit_code = poll_exit_code(task); + if was_running && exit_code.is_some() { + // The child just exited; give the reader threads a moment to flush + // the final output through the pipes before draining the buffer. + std::thread::sleep(std::time::Duration::from_millis(100)); + } + let (new_output, dropped) = drain_task_output(task); + + let command = &task.command; + let status = match exit_code { + None => format!("Task {task_id} (`{command}`) is still running."), + Some(_) if task.killed => format!("Task {task_id} (`{command}`) was killed."), + Some(code) => format!("Task {task_id} (`{command}`) exited with code {code}."), + }; + + if new_output.is_empty() { + format!( + "{status} No new output since the last check.{}", + dropped_note(dropped) + ) + } else { + format!( + "{status} New output since the last check:{}\n{new_output}", + dropped_note(dropped) + ) + } +} + +/// `kill_command` tool: stop a background task and report its output tail. +fn exec_kill_command(arguments: &str) -> String { + let task_id = match parse_task_id(arguments) { + Ok(id) => id, + Err(err) => return err, + }; + + let mut map = lock_tasks(); + let Some(task) = map.get_mut(&task_id) else { + return unknown_task(task_id); + }; + + if let Some(code) = poll_exit_code(task) { + let (new_output, dropped) = drain_task_output(task); + let command = &task.command; + let tail = if new_output.is_empty() { + String::new() + } else { + format!(" Unread output:\n{new_output}") + }; + return format!( + "Task {task_id} (`{command}`) had already exited with code {code}; nothing to \ + kill.{}{tail}", + dropped_note(dropped) + ); + } + + if let Err(err) = task.child.kill() { + return format!("Error: failed to kill task {task_id}: {err}"); + } + match task.child.wait() { + Ok(status) => task.exit_code = Some(status.code().unwrap_or(-1)), + Err(_) => task.exit_code = Some(-1), + } + task.killed = true; + + // Let the reader threads flush whatever was in flight before reporting. + std::thread::sleep(std::time::Duration::from_millis(100)); + let (new_output, dropped) = drain_task_output(task); + + let command = &task.command; + let tail = if new_output.is_empty() { + " It produced no unread output.".to_string() + } else { + format!(" Last output:\n{new_output}") + }; + format!( + "Killed task {task_id} (`{command}`).{}{tail}", + dropped_note(dropped) + ) +} + #[cfg(test)] mod tests { use super::*; @@ -1596,7 +1900,7 @@ mod tests { #[test] fn test_all_tools_count() { let tools = all_tools(); - assert_eq!(tools.len(), 13); + assert_eq!(tools.len(), 15); assert_eq!(tools[0].name, "read_file"); assert_eq!(tools[1].name, "create_directory"); assert_eq!(tools[2].name, "list_directory"); @@ -1610,6 +1914,8 @@ mod tests { assert_eq!(tools[10].name, "glob"); assert_eq!(tools[11].name, "write_todos"); assert_eq!(tools[12].name, "remember"); + assert_eq!(tools[13].name, "command_output"); + assert_eq!(tools[14].name, "kill_command"); } #[test] @@ -2201,4 +2507,192 @@ mod tests { let result = exec_run_command(&args); assert!(result.contains("err"), "got: {result}"); } + + // ── background command tests ───────────────────────────────────────── + + /// Extract the task id from a "Started background task {id}: ..." result. + fn background_task_id(result: &str) -> u64 { + let digits: String = result + .strip_prefix("Started background task ") + .unwrap_or_else(|| panic!("unexpected spawn result: {result}")) + .chars() + .take_while(char::is_ascii_digit) + .collect(); + digits.parse().expect("task id") + } + + /// Whether the task's child has exited, without draining its output. + fn background_task_exited(task_id: u64) -> bool { + let mut map = lock_tasks(); + match map.get_mut(&task_id) { + Some(task) => poll_exit_code(task).is_some(), + None => true, + } + } + + /// Wait (bounded) for a background task's child to exit. + fn wait_for_background_exit(task_id: u64) { + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(30); + while !background_task_exited(task_id) { + assert!( + std::time::Instant::now() < deadline, + "task {task_id} did not exit in time" + ); + std::thread::sleep(std::time::Duration::from_millis(50)); + } + } + + #[test] + fn test_run_command_background_lifecycle() { + #[cfg(unix)] + let command = "echo start; sleep 2; echo done"; + #[cfg(windows)] + let command = "echo start&& ping -n 3 127.0.0.1 > nul&& echo done"; + + let args = serde_json::json!({ + "command": command, + "cwd": std::env::temp_dir(), + "run_in_background": true + }) + .to_string(); + let spawned = std::time::Instant::now(); + let result = exec_run_command(&args); + // Spawning must return immediately, not wait the ~2s the command takes. + assert!( + spawned.elapsed() < std::time::Duration::from_secs(1), + "background spawn blocked for {:?}", + spawned.elapsed() + ); + assert!(result.contains("command_output"), "got: {result}"); + let task_id = background_task_id(&result); + + // Polling while the command is still sleeping reports it as running. + let poll_args = serde_json::json!({ "task_id": task_id }).to_string(); + let poll = exec_command_output(&poll_args); + assert!(poll.contains("still running"), "got: {poll}"); + let mut combined = poll; + + wait_for_background_exit(task_id); + // Grace period so the reader threads finish draining the pipes. + std::thread::sleep(std::time::Duration::from_millis(300)); + let final_poll = exec_command_output(&poll_args); + assert!( + final_poll.contains("exited with code 0"), + "got: {final_poll}" + ); + combined.push_str(&final_poll); + + // Across the polls, all output the command printed was delivered. + assert!(combined.contains("start"), "got: {combined}"); + assert!(combined.contains("done"), "got: {combined}"); + } + + #[test] + fn test_kill_command_stops_background_task() { + #[cfg(unix)] + let command = "sleep 30"; + #[cfg(windows)] + let command = "ping -n 31 127.0.0.1 > nul"; + + let args = serde_json::json!({ + "command": command, + "cwd": std::env::temp_dir(), + "run_in_background": true + }) + .to_string(); + let result = exec_run_command(&args); + let task_id = background_task_id(&result); + + let kill_args = serde_json::json!({ "task_id": task_id }).to_string(); + let killed_at = std::time::Instant::now(); + let kill_result = exec_kill_command(&kill_args); + assert!( + kill_result.contains(&format!("Killed task {task_id}")), + "got: {kill_result}" + ); + // The kill must not wait out the 30s sleep. + assert!( + killed_at.elapsed() < std::time::Duration::from_secs(5), + "kill blocked for {:?}", + killed_at.elapsed() + ); + + // A later poll reports the task as killed, not still running. + let poll = exec_command_output(&serde_json::json!({ "task_id": task_id }).to_string()); + assert!(poll.contains("was killed"), "got: {poll}"); + } + + #[test] + fn test_command_output_unknown_task() { + let result = exec_command_output(r#"{"task_id": 9999999}"#); + assert!( + result.contains("no background task with id"), + "got: {result}" + ); + } + + #[test] + fn test_kill_command_unknown_task() { + let result = exec_kill_command(r#"{"task_id": 9999999}"#); + assert!( + result.contains("no background task with id"), + "got: {result}" + ); + } + + #[test] + fn test_command_output_missing_task_id() { + let result = exec_command_output("{}"); + assert!( + result.contains("missing required parameter"), + "got: {result}" + ); + } + + #[test] + fn test_background_output_is_capped() { + // Print well over COMMAND_OUTPUT_LIMIT (50 000) bytes without polling, + // so the buffer must drop the oldest output and note the truncation. + #[cfg(unix)] + let command = r#"i=0; while [ $i -lt 200 ]; do printf 'line-%04d %s\n' "$i" \ + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa \ + aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa \ + aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa \ + aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa \ + aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; \ + i=$((i+1)); done"#; + #[cfg(windows)] + let command = "powershell -NoProfile -Command \"1..200 | ForEach-Object { \ + ('line-{0:d4} ' -f $_) + ('a' * 330) }\""; + + let args = serde_json::json!({ + "command": command, + "cwd": std::env::temp_dir(), + "run_in_background": true + }) + .to_string(); + let result = exec_run_command(&args); + let task_id = background_task_id(&result); + + wait_for_background_exit(task_id); + // Grace period so the reader threads finish draining the pipes. + std::thread::sleep(std::time::Duration::from_millis(300)); + + let poll = exec_command_output(&serde_json::json!({ "task_id": task_id }).to_string()); + assert!(poll.contains("exited with code 0"), "got: {poll}"); + assert!( + poll.contains("earlier output was dropped"), + "expected truncation note, got: {poll}" + ); + // The oldest lines were dropped; the newest survived. + assert!(!poll.contains("line-0000"), "oldest output not dropped"); + assert!(poll.contains("line-0199"), "newest output missing"); + // Buffer stayed within the cap, plus the framing: the status line + // (which quotes the command) and the truncation note. + assert!( + poll.len() <= COMMAND_OUTPUT_LIMIT + command.len() + 500, + "poll result too large: {} bytes", + poll.len() + ); + } }