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: fix wait_for_task/run_task todo race with a refcounted waiter registry

Browse source
This commit is contained in:
damocles 2026-08-11 18:14:57 +02:00 • committed by mara
commit fc2f75fed1
1 changed files with 140 additions and 96 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -72,38 +72,79 @@ fn running() -> &'static Mutex<HashMap<String, RunningHandle>> {
}
// ---------------------------------------------------------------------------
// Wake suppression: skip the completion wake when a caller already
// synchronously observed the task's terminal state via an inline
// `wait_seconds` poll on `BashRun` or `BashStatus` — the tool response
// already delivered the result in that same turn, so a follow-up wake
// message would just be a redundant duplicate of information the agent has.
// Inline-waiter registry: skip the completion todo when a caller is
// currently (or was, until moments ago) synchronously polling the task's
// terminal state via `wait_seconds` on `BashRun` or `BashStatus` — the tool
// response already delivers the result in that same turn, so a follow-up
// todo would just be a redundant duplicate of information the agent has.
//
// Superseded a one-shot flag set by the waiter *after* observing terminal
// state and checked by the runner on its own independent poll schedule —
// two independently timed reads of *different* state (task file vs. flag)
// can't be made race-free by reordering, only by sharing a lock. This
// shape does: both sides check/mutate **one registry** under **one lock**,
// deciding on presence rather than a flag that might be set too late. A
// still-registered waiter's own poll loop is guaranteed to see the
// just-written terminal file on its next iteration — same file, not a
// message that could be missed — so skipping its todo can't strand the
// agent. See the forge issue tracker for the full before/after trace.
// ---------------------------------------------------------------------------
/// In-memory only — a daemon restart wipes it, which is fine: a task still
/// `running` across a restart is marked `interrupted` on boot (see module
/// docs) and gets its own fresh wake, independent of this set.
fn wake_suppressed() -> &'static Mutex<HashSet<String>> {
static SUPPRESSED: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();
SUPPRESSED.get_or_init(|| Mutex::new(HashSet::new()))
/// Refcounted, not a plain set: two concurrent inline waiters on the same
/// `id` are possible (e.g. `run` and a separate `status` call racing each
/// other), and a plain `HashSet` would let the *first* one to drop
/// deregister the id out from under the second, still-live one. Counting
/// means presence only goes to zero once every registered guard has
/// dropped. In-memory only — a daemon restart wipes it, which is fine: any
/// waiter registered here belonged to an in-flight MCP call the restart
/// also tore down, and a task still `running` across a restart is marked
/// `interrupted` on boot (see module docs) and gets its own fresh todo,
/// independent of this map.
fn waiting_ids() -> &'static Mutex<HashMap<String, u32>> {
static WAITING: OnceLock<Mutex<HashMap<String, u32>>> = OnceLock::new();
WAITING.get_or_init(|| Mutex::new(HashMap::new()))
}
/// Mark `id`'s completion wake as already-delivered-inline. Called after an
/// inline `wait_seconds` poll (on `BashRun` or `BashStatus`) observes a
/// terminal task, before the response carrying the full status is written
/// back to the caller. Idempotent — safe to call more than once per id.
pub(crate) fn suppress_wake(id: &str) {
wake_suppressed().lock().unwrap().insert(id.to_owned());
/// True if an inline waiter is currently registered for `id`. Checked by
/// `run_task`'s completion handler under the same lock [`WaiterGuard`]
/// (de)registers under — "is a waiter here" and "a waiter leaving" can
/// never observe torn state relative to each other, unlike the flag this
/// replaced.
fn waiter_present(id: &str) -> bool {
waiting_ids().lock().unwrap().contains_key(id)
}
/// Consume (remove + report) `id`'s suppression flag. Returns `true` if the
/// wake should be skipped. One-shot by intent: a completed task's entry is
/// meant to be drained exactly once. There's a known low-probability race
/// (see `run_task`'s completion handler) where `suppress_wake` for a task
/// fires *after* this has already run for it — if the same name gets reused
/// before that late insert lands, the dangling flag could suppress the new
/// task's wake instead. Not eliminated, just narrow.
fn take_wake_suppressed(id: &str) -> bool {
wake_suppressed().lock().unwrap().remove(id)
/// RAII registration for one `wait_for_task` call. Increments `id`'s count
/// on construction, decrements (removing the entry at zero) on drop —
/// covers every `wait_for_task` return path (terminal found, deadline hit,
/// task vanished) with one code path instead of duplicating the removal at
/// each `return`. Rust drops locals after the return expression is
/// evaluated, so the guard stays registered through `wait_for_task`'s very
/// last read of the task file and only deregisters right before the value
/// actually returns to the caller.
struct WaiterGuard<'a>(&'a str);
impl<'a> WaiterGuard<'a> {
fn new(id: &'a str) -> Self {
*waiting_ids()
.lock()
.unwrap()
.entry(id.to_owned())
.or_insert(0) += 1;
Self(id)
}
}
impl Drop for WaiterGuard<'_> {
fn drop(&mut self) {
let mut waiting = waiting_ids().lock().unwrap();
if let Some(count) = waiting.get_mut(self.0) {
*count -= 1;
if *count == 0 {
waiting.remove(self.0);
}
}
}
}
/// Outcome of one `exec_cmd` run.
@ -287,54 +328,46 @@ pub fn submit_task(cmd: String, timeout_secs: Option<u64>, name: Option<String>)
/// Inline wait: poll `read_task(id)` until terminal state or deadline.
/// Returns the final task on success, or `None` if it never completed.
///
/// Whenever a terminal task is observed here, the caller is about to receive
/// that result directly in its tool response — so the completion wake for
/// `id` is marked [`suppress_wake`]d: a status query (waited or not) that
/// already told the agent the outcome shouldn't be followed by a redundant
/// "task finished" inbox message for the same information.
/// Registered as an inline waiter for `id` (via [`WaiterGuard`]) for the
/// polling loop only — deliberately **not** across the final fallback read
/// below. A guard still held during that last read would let
/// `run_task`'s completion handler see "waiter present" and skip the
/// todo for a result this call already fixed (non-terminal, from the read
/// that just lost the deadline race) — a genuinely dropped notification,
/// not a redundant one. Dropping the guard first means the worst case if
/// `run_task` completes in the gap between the guard dropping and this
/// read running is a harmless extra todo (the pre-existing, documented
/// tolerance), never a missed one.
pub async fn wait_for_task(id: &str, wait_secs: u64) -> Option<TaskFile> {
let cap = wait_secs.min(MAX_WAIT_SECS);
if cap == 0 {
return observe_terminal(read_task(id));
return read_task(id);
}
let deadline = tokio::time::Instant::now() + Duration::from_secs(cap);
loop {
match read_task(id) {
None => break,
Some(task) => {
if matches!(
task.status,
TaskStatus::Done
| TaskStatus::TimedOut
| TaskStatus::Interrupted
| TaskStatus::Killed
) {
suppress_wake(id);
return Some(task);
{
let _waiting = WaiterGuard::new(id);
loop {
match read_task(id) {
None => break,
Some(task) => {
if matches!(
task.status,
TaskStatus::Done
| TaskStatus::TimedOut
| TaskStatus::Interrupted
| TaskStatus::Killed
) {
return Some(task);
}
}
}
if tokio::time::Instant::now() >= deadline {
break;
}
tokio::time::sleep(Duration::from_millis(POLL_MS)).await;
}
if tokio::time::Instant::now() >= deadline {
break;
}
tokio::time::sleep(Duration::from_millis(POLL_MS)).await;
}
observe_terminal(read_task(id))
}
/// Marks the wake suppressed if `task` is in a terminal state; passes
/// `task` through unchanged either way. Shared tail helper for both
/// `wait_for_task` return points.
fn observe_terminal(task: Option<TaskFile>) -> Option<TaskFile> {
if let Some(t) = &task
&& matches!(
t.status,
TaskStatus::Done | TaskStatus::TimedOut | TaskStatus::Interrupted | TaskStatus::Killed
)
{
suppress_wake(&t.id);
}
task
} // guard dropped before the final read below
read_task(id)
}
/// Kill a running or still-pending task.
@ -565,17 +598,18 @@ async fn run_task(mut task: TaskFile, socket: &Path) {
if let Err(e) = write_task(&task) {
tracing::warn!(id = %id, error = ?e, "bash_runner: write done state failed");
}
// If an inline `status`/`run` wait already handed this exact terminal
// result to the caller in a tool response (see `wait_for_task` /
// `observe_terminal`), there's nothing left to surface — retire the
// keyed todo without a `done` upsert so the agent gets no redundant
// wake. Narrow race: an inline waiter polling in the few hundred ms
// around this point may lose the race and still get a `done` todo
// alongside its inline result; best-effort, same tolerance as the rest
// of this daemon's guarantees.
if take_wake_suppressed(&id) {
// If an inline `status`/`run` wait is currently registered for this
// task (see `WaiterGuard`), that call's own poll loop is guaranteed to
// observe the terminal file just written above on its very next
// iteration — same file, not a message that could be missed — so
// there's nothing left for a todo to surface: retire the keyed todo
// without a `done` upsert instead. `waiter_present` and
// `WaiterGuard`'s (de)registration share one lock, so this can't race
// the way the old flag-based check could (see the registry's module
// doc for the full before/after).
if waiter_present(&id) {
clear_bash_todo(socket, &id).await;
tracing::debug!(id = %id, "bash_runner: done todo suppressed (already observed via status)");
tracing::debug!(id = %id, "bash_runner: done todo suppressed (inline waiter registered)");
return;
}
@ -833,31 +867,41 @@ mod tests {
assert!(validate_task_name(&"x".repeat(Ident::MAX_LEN)).is_ok());
}
// Wake suppression is a single process-wide registry (see
// `wake_suppressed()`), so these run serially against distinct ids to
// avoid cross-test interference under parallel test execution.
// The waiter registry is a single process-wide `Mutex<HashMap<..>>`
// refcount (see `waiting_ids()`), so these run serially against
// distinct ids to avoid cross-test interference under parallel test
// execution.
#[test]
fn wake_suppression_is_one_shot() {
use super::{suppress_wake, take_wake_suppressed};
let id = "test-2270-one-shot";
assert!(!take_wake_suppressed(id), "unset id starts unsuppressed");
suppress_wake(id);
assert!(take_wake_suppressed(id), "set id reports suppressed once");
assert!(
!take_wake_suppressed(id),
"consuming the flag clears it — second read is unsuppressed"
);
fn waiter_guard_registers_and_deregisters_on_drop() {
use super::{WaiterGuard, waiter_present};
let id = "test-2905-guard-lifecycle";
assert!(!waiter_present(id), "unregistered id starts absent");
{
let _guard = WaiterGuard::new(id);
assert!(waiter_present(id), "present for the guard's lifetime");
}
assert!(!waiter_present(id), "guard drop deregisters");
}
#[test]
fn wake_suppression_is_idempotent_to_set() {
use super::{suppress_wake, take_wake_suppressed};
let id = "test-2270-idempotent";
suppress_wake(id);
suppress_wake(id); // simulates two concurrent observers of the same terminal task
assert!(take_wake_suppressed(id));
assert!(!take_wake_suppressed(id));
fn waiter_guard_stays_present_until_every_sibling_drops() {
use super::{WaiterGuard, waiter_present};
let id = "test-2905-guard-refcount";
let first = WaiterGuard::new(id);
let second = WaiterGuard::new(id); // simulates two concurrent inline waiters
assert!(waiter_present(id));
drop(first);
// Refcounted, not a plain set: a sibling guard's early drop must
// not un-register an id a still-live guard needs — `second` is
// still registered, so presence must survive `first` alone
// dropping.
assert!(
waiter_present(id),
"still present while a sibling guard is live"
);
drop(second);
assert!(!waiter_present(id), "absent once every guard has dropped");
}
// done_summary: the stderr-present branch must trigger on has_stderr