From fc2f75fed1b46ecbea95e81fccf358d532577805 Mon Sep 17 00:00:00 2001 From: damocles Date: Tue, 11 Aug 2026 18:14:57 +0200 Subject: [PATCH] hive-bash-mcp: fix wait_for_task/run_task todo race with a refcounted waiter registry --- hive-bash-mcp/src/runner.rs | 236 +++++++++++++++++++++--------------- 1 file changed, 140 insertions(+), 96 deletions(-) diff --git a/hive-bash-mcp/src/runner.rs b/hive-bash-mcp/src/runner.rs index 7fa0c975..16423582 100644 --- a/hive-bash-mcp/src/runner.rs +++ b/hive-bash-mcp/src/runner.rs @@ -72,38 +72,79 @@ fn running() -> &'static Mutex> { } // --------------------------------------------------------------------------- -// 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> { - static SUPPRESSED: OnceLock>> = 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> { + static WAITING: OnceLock>> = 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, name: Option) /// 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 { 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) -> Option { - 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>` + // 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