hive-bash-mcp: fix wait_for_task/run_task todo race with a refcounted waiter registry

hive-bash-mcp/src/runner.rs

@@ -72,38 +72,79 @@ fn running() -> &'static Mutex<HashMap<String, RunningHandle>> {
 }
 
 // ---------------------------------------------------------------------------
-// Wake suppression: skip the completion wake when a caller already
+// Inline-waiter registry: skip the completion todo when a caller is
-// synchronously observed the task's terminal state via an inline
+// currently (or was, until moments ago) synchronously polling the task's
-// `wait_seconds` poll on `BashRun` or `BashStatus` — the tool response
+// terminal state via `wait_seconds` on `BashRun` or `BashStatus` — the tool
-// already delivered the result in that same turn, so a follow-up wake
+// response already delivers the result in that same turn, so a follow-up
-// message would just be a redundant duplicate of information the agent has.
+// 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
+/// Refcounted, not a plain set: two concurrent inline waiters on the same
-/// `running` across a restart is marked `interrupted` on boot (see module
+/// `id` are possible (e.g. `run` and a separate `status` call racing each
-/// docs) and gets its own fresh wake, independent of this set.
+/// other), and a plain `HashSet` would let the *first* one to drop
-fn wake_suppressed() -> &'static Mutex<HashSet<String>> {
+/// deregister the id out from under the second, still-live one. Counting
-    static SUPPRESSED: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();
+/// means presence only goes to zero once every registered guard has
-    SUPPRESSED.get_or_init(|| Mutex::new(HashSet::new()))
+/// 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
+/// True if an inline waiter is currently registered for `id`. Checked by
-/// inline `wait_seconds` poll (on `BashRun` or `BashStatus`) observes a
+/// `run_task`'s completion handler under the same lock [`WaiterGuard`]
-/// terminal task, before the response carrying the full status is written
+/// (de)registers under — "is a waiter here" and "a waiter leaving" can
-/// back to the caller. Idempotent — safe to call more than once per id.
+/// never observe torn state relative to each other, unlike the flag this
-pub(crate) fn suppress_wake(id: &str) {
+/// replaced.
-    wake_suppressed().lock().unwrap().insert(id.to_owned());
+fn waiter_present(id: &str) -> bool {
+    waiting_ids().lock().unwrap().contains_key(id)
 }
 
-/// Consume (remove + report) `id`'s suppression flag. Returns `true` if the
+/// RAII registration for one `wait_for_task` call. Increments `id`'s count
-/// wake should be skipped. One-shot by intent: a completed task's entry is
+/// on construction, decrements (removing the entry at zero) on drop —
-/// meant to be drained exactly once. There's a known low-probability race
+/// covers every `wait_for_task` return path (terminal found, deadline hit,
-/// (see `run_task`'s completion handler) where `suppress_wake` for a task
+/// task vanished) with one code path instead of duplicating the removal at
-/// fires *after* this has already run for it — if the same name gets reused
+/// each `return`. Rust drops locals after the return expression is
-/// before that late insert lands, the dangling flag could suppress the new
+/// evaluated, so the guard stays registered through `wait_for_task`'s very
-/// task's wake instead. Not eliminated, just narrow.
+/// last read of the task file and only deregisters right before the value
-fn take_wake_suppressed(id: &str) -> bool {
+/// actually returns to the caller.
-    wake_suppressed().lock().unwrap().remove(id)
+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
+/// Registered as an inline waiter for `id` (via [`WaiterGuard`]) for the
-/// that result directly in its tool response — so the completion wake for
+/// polling loop only — deliberately **not** across the final fallback read
-/// `id` is marked [`suppress_wake`]d: a status query (waited or not) that
+/// below. A guard still held during that last read would let
-/// already told the agent the outcome shouldn't be followed by a redundant
+/// `run_task`'s completion handler see "waiter present" and skip the
-/// "task finished" inbox message for the same information.
+/// 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) {
+        let _waiting = WaiterGuard::new(id);
-            None => break,
+        loop {
-            Some(task) => {
+            match read_task(id) {
-                if matches!(
+                None => break,
-                    task.status,
+                Some(task) => {
-                    TaskStatus::Done
+                    if matches!(
-                        | TaskStatus::TimedOut
+                        task.status,
-                        | TaskStatus::Interrupted
+                        TaskStatus::Done
-                        | TaskStatus::Killed
+                            | TaskStatus::TimedOut
-                ) {
+                            | TaskStatus::Interrupted
-                    suppress_wake(id);
+                            | TaskStatus::Killed
-                    return Some(task);
+                    ) {
+                        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 {
+    } // guard dropped before the final read below
-            break;
+    read_task(id)
-        }
-        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
 }
 
 /// 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
+    // If an inline `status`/`run` wait is currently registered for this
-    // result to the caller in a tool response (see `wait_for_task` /
+    // task (see `WaiterGuard`), that call's own poll loop is guaranteed to
-    // `observe_terminal`), there's nothing left to surface — retire the
+    // observe the terminal file just written above on its very next
-    // keyed todo without a `done` upsert so the agent gets no redundant
+    // iteration — same file, not a message that could be missed — so
-    // wake. Narrow race: an inline waiter polling in the few hundred ms
+    // there's nothing left for a todo to surface: retire the keyed todo
-    // around this point may lose the race and still get a `done` todo
+    // without a `done` upsert instead. `waiter_present` and
-    // alongside its inline result; best-effort, same tolerance as the rest
+    // `WaiterGuard`'s (de)registration share one lock, so this can't race
-    // of this daemon's guarantees.
+    // the way the old flag-based check could (see the registry's module
-    if take_wake_suppressed(&id) {
+    // 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
+    // The waiter registry is a single process-wide `Mutex<HashMap<..>>`
-    // `wake_suppressed()`), so these run serially against distinct ids to
+    // refcount (see `waiting_ids()`), so these run serially against
-    // avoid cross-test interference under parallel test execution.
+    // distinct ids to avoid cross-test interference under parallel test
+    // execution.
 
     #[test]
-    fn wake_suppression_is_one_shot() {
+    fn waiter_guard_registers_and_deregisters_on_drop() {
-        use super::{suppress_wake, take_wake_suppressed};
+        use super::{WaiterGuard, waiter_present};
-        let id = "test-2270-one-shot";
+        let id = "test-2905-guard-lifecycle";
-        assert!(!take_wake_suppressed(id), "unset id starts unsuppressed");
+        assert!(!waiter_present(id), "unregistered id starts absent");
-        suppress_wake(id);
+        {
-        assert!(take_wake_suppressed(id), "set id reports suppressed once");
+            let _guard = WaiterGuard::new(id);
-        assert!(
+            assert!(waiter_present(id), "present for the guard's lifetime");
-            !take_wake_suppressed(id),
+        }
-            "consuming the flag clears it — second read is unsuppressed"
+        assert!(!waiter_present(id), "guard drop deregisters");
-        );
     }
 
     #[test]
-    fn wake_suppression_is_idempotent_to_set() {
+    fn waiter_guard_stays_present_until_every_sibling_drops() {
-        use super::{suppress_wake, take_wake_suppressed};
+        use super::{WaiterGuard, waiter_present};
-        let id = "test-2270-idempotent";
+        let id = "test-2905-guard-refcount";
-        suppress_wake(id);
+        let first = WaiterGuard::new(id);
-        suppress_wake(id); // simulates two concurrent observers of the same terminal task
+        let second = WaiterGuard::new(id); // simulates two concurrent inline waiters
-        assert!(take_wake_suppressed(id));
+        assert!(waiter_present(id));
-        assert!(!take_wake_suppressed(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