vinzenz/hyperhive
vinzenz/hyperhive
Watch
0
0
Fork
You've already forked hyperhive
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

hive-bash-mcp: add kill tool to stop a running or pending bash task

Browse source
This commit is contained in:
damocles 2026-06-19 01:34:18 +02:00 • committed by mara
commit b16629801b
4 changed files with 291 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

253
hive-bash-mcp/src/runner.rs
View file

@ -14,14 +14,15 @@
//! The runner kills the child process on timeout — `tokio::process::Child::drop()`
//! does not kill children, so we explicitly call `child.kill().await`.
use std::collections::HashSet;
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use anyhow::{Result, bail};
use tokio::io::AsyncWriteExt as _;
use tokio::sync::Notify;
use crate::paths;
use crate::protocol::{TaskFile, TaskStatus};
@ -41,6 +42,50 @@ const POLL_MS: u64 = 100;
static TASK_SEQ: AtomicU64 = AtomicU64::new(0);
// ---------------------------------------------------------------------------
// Running-task registry (for on-demand kill)
// ---------------------------------------------------------------------------
/// Handle to a currently-executing task, used by [`kill_task`] to signal it.
/// `exec_cmd` selects on `cancel`; when notified it signals the task's
/// process group with `SIGKILL` (if `force`) or `SIGINT`.
struct RunningHandle {
cancel: Arc<Notify>,
force: Arc<AtomicBool>,
}
/// Global registry of running tasks (id → handle). Populated by `run_task`
/// for the duration of execution and read by [`kill_task`] from the socket
/// dispatch path — a separate async context from the runner loop, so a
/// shared global (rather than the loop-local `claimed` set) is needed.
fn running() -> &'static Mutex<HashMap<String, RunningHandle>> {
static RUNNING: OnceLock<Mutex<HashMap<String, RunningHandle>>> = OnceLock::new();
RUNNING.get_or_init(|| Mutex::new(HashMap::new()))
}
/// Outcome of one `exec_cmd` run.
enum ExecOutcome {
/// Process exited on its own with this status code.
Exited(i32),
/// Killed by the per-task timeout (`timeout_secs`).
TimedOut,
/// Killed on request via [`kill_task`]. `forced` = SIGKILL vs SIGINT.
Killed { forced: bool },
}
/// Send `sig` to the task's process group. A negative pgid targets the whole
/// group, so `bash -c` plus any children it spawned all receive the signal.
/// No-op if the child had no pid (already reaped).
fn signal_group(pgid: Option<i32>, sig: i32) {
if let Some(pgid) = pgid {
// SAFETY: plain libc `kill(2)`; `-pgid` targets the child's own
// process group (it called `setpgid(0, 0)` so it leads its group).
unsafe {
libc::kill(-pgid, sig);
}
}
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
@ -232,7 +277,10 @@ pub async fn wait_for_task(id: &str, wait_secs: u64) -> Option<TaskFile> {
Some(task) => {
if matches!(
task.status,
TaskStatus::Done | TaskStatus::TimedOut | TaskStatus::Interrupted
TaskStatus::Done
| TaskStatus::TimedOut
| TaskStatus::Interrupted
| TaskStatus::Killed
) {
return Some(task);
}
@ -246,6 +294,42 @@ pub async fn wait_for_task(id: &str, wait_secs: u64) -> Option<TaskFile> {
read_task(id)
}
/// Kill a running or still-pending task.
///
/// - **Running** (in the registry): signals its process group via the
/// `exec_cmd` cancel path — `SIGKILL` if `force`, else `SIGINT`. The task
/// transitions to [`TaskStatus::Killed`] once the process exits and a
/// completion wake fires as usual. A `SIGINT` relies on the process
/// honouring it; pass `force` for a guaranteed stop.
/// - **Pending** (queued, not yet started): marked `Killed` directly so the
/// runner loop never starts it. No process exists yet, so `force` is moot.
/// - **Terminal or unknown id**: no-op.
///
/// Returns `(killed, was_running)`: `killed` = a kill was issued (signal sent
/// or pending task cancelled); `was_running` = the task was actively running.
pub fn kill_task(id: &str, force: bool) -> (bool, bool) {
// Running task: signal via the registry.
{
let guard = running().lock().unwrap();
if let Some(handle) = guard.get(id) {
handle.force.store(force, Ordering::SeqCst);
handle.cancel.notify_one();
return (true, true);
}
}
// Pending task: mark Killed so the poll loop never starts it.
if let Some(mut task) = read_task(id)
&& task.status == TaskStatus::Pending
{
task.status = TaskStatus::Killed;
task.completed_at = Some(now_unix());
let _ = write_task(&task);
refresh_loose_ends();
return (true, false);
}
(false, false)
}
// ---------------------------------------------------------------------------
// Runner background loop
// ---------------------------------------------------------------------------
@ -358,27 +442,57 @@ async fn run_task(mut task: TaskFile, socket: &Path) {
let out_path = paths::task_out(&id);
let err_path = paths::task_err(&id);
let (timed_out, exit_code) =
match exec_cmd(&task.cmd, &out_path, &err_path, task.timeout_secs).await {
Ok((code, false)) => (false, Some(code)),
Ok((_, true)) => {
tracing::warn!(id = %id, "bash_runner: task timed out");
(true, None)
}
Err(e) => {
tracing::warn!(id = %id, error = ?e, "bash_runner: exec error");
(false, None)
}
};
// Register a kill handle for the duration of execution so `kill_task`
// (called from the socket dispatch path) can signal this task.
let cancel = Arc::new(Notify::new());
let force = Arc::new(AtomicBool::new(false));
running().lock().unwrap().insert(
id.clone(),
RunningHandle {
cancel: cancel.clone(),
force: force.clone(),
},
);
let outcome = exec_cmd(
&task.cmd,
&out_path,
&err_path,
task.timeout_secs,
&cancel,
&force,
)
.await;
// Deregister before writing terminal state — the task is no longer killable.
running().lock().unwrap().remove(&id);
let stdout_tail = tail_file(&out_path, SUMMARY_BYTES);
let stderr_tail = tail_file(&err_path, SUMMARY_BYTES);
task.status = if timed_out {
TaskStatus::TimedOut
} else {
TaskStatus::Done
let (status, exit_code, summary) = match outcome {
Ok(ExecOutcome::Exited(code)) => (TaskStatus::Done, Some(code), format!("exit={code}")),
Ok(ExecOutcome::TimedOut) => {
tracing::warn!(id = %id, "bash_runner: task timed out");
let secs = task.timeout_secs.unwrap_or(0);
(
TaskStatus::TimedOut,
None,
format!("timed out after {secs}s"),
)
}
Ok(ExecOutcome::Killed { forced }) => {
let sig = if forced { "SIGKILL" } else { "SIGINT" };
tracing::warn!(id = %id, sig, "bash_runner: task killed");
(TaskStatus::Killed, None, format!("killed ({sig})"))
}
Err(e) => {
tracing::warn!(id = %id, error = ?e, "bash_runner: exec error");
(TaskStatus::Done, None, "exec error".to_owned())
}
};
task.status = status;
task.completed_at = Some(now_unix());
task.exit_code = exit_code;
task.stdout_tail = stdout_tail.clone().filter(|s| !s.is_empty());
@ -389,33 +503,31 @@ async fn run_task(mut task: TaskFile, socket: &Path) {
}
refresh_loose_ends();
let summary = if timed_out {
let secs = task.timeout_secs.unwrap_or(0);
format!("timed out after {secs}s")
} else {
format!("exit={}", exit_code.unwrap_or(-1))
};
let out_snippet = stdout_tail.as_deref().unwrap_or("").trim();
let err_snippet = stderr_tail.as_deref().unwrap_or("").trim();
send_wake(socket, &id, &summary, Some((out_snippet, err_snippet))).await;
}
/// Run `bash -c cmd`, streaming output to files. Returns `(exit_code, timed_out)`.
/// Run `bash -c cmd` in its own process group, streaming output to files.
/// `timeout_secs = None` means no timeout — run until natural exit.
///
/// Tasks run under `bash`, not `sh`: on NixOS `/bin/sh` is bash in POSIX
/// mode, which disables bashisms (arrays, `[[ … ]]`, `local`, process
/// substitution, …). Agents write bash, so we invoke `bash` (via
/// `/usr/bin/env bash`).
/// Tasks run under `bash`, not `sh` (on NixOS `/bin/sh` is bash in POSIX
/// mode, which disables bashisms — arrays, `[[ … ]]`, `local`, process
/// substitution), invoked via `/usr/bin/env bash`.
///
/// The task is cancellable via `cancel` (set `force` first): on cancel the
/// whole process group is signalled — `SIGKILL` if `force`, else `SIGINT` —
/// so children of the shell die too, not just the shell. The timeout path
/// likewise `SIGKILL`s the group.
async fn exec_cmd(
cmd: &str,
out_path: &Path,
err_path: &Path,
timeout_secs: Option<u64>,
) -> Result<(i32, bool)> {
cancel: &Notify,
force: &AtomicBool,
) -> Result<ExecOutcome> {
use tokio::process::Command;
// SAFETY: `nice` is async-signal-safe and modifies only the calling
// process's scheduling priority before exec. No allocations, no locks.
// `/usr/bin/env bash` rather than a bare `bash`: `/usr/bin/env` is at a
// fixed absolute path (coreutils, present on NixOS), and it resolves
// `bash` via PATH — the same controlled PATH the daemon's systemd unit
@ -427,13 +539,20 @@ async fn exec_cmd(
.arg(cmd)
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
// SAFETY: `setpgid`/`nice` are async-signal-safe and only touch the
// calling (child) process before exec — no allocations, no locks.
// `setpgid(0, 0)` makes the child its own process-group leader so the
// whole subtree can be signalled as a group (`kill(-pgid, …)`).
unsafe {
cmd_builder.pre_exec(|| {
libc::setpgid(0, 0);
libc::nice(10);
Ok(())
});
}
let mut child = cmd_builder.spawn()?;
// With `setpgid(0, 0)` the child's pgid equals its pid.
let pgid = child.id().map(u32_to_i32);
let stdout = child.stdout.take().expect("stdout piped");
let stderr = child.stderr.take().expect("stderr piped");
@ -450,33 +569,53 @@ async fn exec_cmd(
err_path,
));
if let Some(secs) = timeout_secs {
match tokio::time::timeout(Duration::from_secs(secs), child.wait()).await {
Ok(Ok(status)) => {
let _ = copy_out.await;
let _ = copy_err.await;
Ok((status.code().unwrap_or(-1), false))
// A future that fires when the timeout elapses, or never if `None`.
let timeout_fut = async {
match timeout_secs {
Some(secs) => tokio::time::sleep(Duration::from_secs(secs)).await,
None => std::future::pending::<()>().await,
}
};
let mut timed_out = false;
let mut killed: Option<bool> = None;
let wait_res = {
let wait = child.wait();
tokio::pin!(wait);
tokio::pin!(timeout_fut);
tokio::select! {
res = &mut wait => res,
() = &mut timeout_fut => {
signal_group(pgid, libc::SIGKILL);
timed_out = true;
(&mut wait).await
}
Ok(Err(e)) => Err(e.into()),
Err(_elapsed) => {
let _ = child.kill().await;
let _ = child.wait().await;
let _ = copy_out.await;
let _ = copy_err.await;
Ok((-1, true))
() = cancel.notified() => {
let forced = force.load(Ordering::SeqCst);
signal_group(pgid, if forced { libc::SIGKILL } else { libc::SIGINT });
killed = Some(forced);
(&mut wait).await
}
}
};
let _ = copy_out.await;
let _ = copy_err.await;
let status = wait_res?;
let outcome = if let Some(forced) = killed {
ExecOutcome::Killed { forced }
} else if timed_out {
ExecOutcome::TimedOut
} else {
// No timeout — wait for natural exit.
match child.wait().await {
Ok(status) => {
let _ = copy_out.await;
let _ = copy_err.await;
Ok((status.code().unwrap_or(-1), false))
}
Err(e) => Err(e.into()),
}
}
ExecOutcome::Exited(status.code().unwrap_or(-1))
};
Ok(outcome)
}
/// `u32` pid → `i32` for `kill(2)`. Pids fit in `i32`; saturates defensively.
fn u32_to_i32(v: u32) -> i32 {
i32::try_from(v).unwrap_or(i32::MAX)
}
async fn copy_stream_to_file<R>(mut reader: R, path: PathBuf)