subagent: close the start/continue TOCTOU race with an atomic reservation

hive-subagent-mcp/src/mcp.rs

@@ -1,5 +1,6 @@
-//! The MCP tool surface: `start` / `continue` / `interrupt`, served
-//! directly over streamable-http — no stdio bridge, no round-trip socket.
+//! The MCP tool surface: `start` / `continue` / `status` / `interrupt`,
+//! served directly over streamable-http — no stdio bridge, no round-trip
+//! socket.
 
 use std::sync::Arc;
 
@@ -116,9 +117,11 @@ impl SubagentMcp {
     }
 
     #[tool(
-        description = "Signal a currently-running subagent session to stop. Only works while \
-            it's actually running — there's no queued/pending state to cancel pre-emptively, \
-            only running or not tracked at all. `force: true` for SIGKILL, otherwise SIGINT."
+        description = "Signal a currently-running subagent session to stop. Only works once \
+            it's actually running — a `start`/`continue` still in its brief window before the \
+            process is confirmed spawned refuses interrupt too (nothing to signal yet; retry \
+            shortly), same as a name with nothing tracked at all. `force: true` for SIGKILL, \
+            otherwise SIGINT."
     )]
     fn interrupt(&self, Parameters(args): Parameters<InterruptArgs>) -> String {
         match session::interrupt(&self.state, &args.name, args.force) {
@@ -129,9 +132,10 @@ impl SubagentMcp {
 
     #[tool(
         description = "Report whether a subagent is currently running — a zero-cost check that \
-            never launches a process, unlike `continue`. Distinguishes running, idle (a session \
-            exists but nothing is in flight — `continue` to give it another turn), and no such \
-            session at all."
+            never launches a process, unlike `continue`. Distinguishes running, starting (a \
+            `start`/`continue` is in flight but not yet a confirmed spawn — this is normally \
+            over in well under a second), idle (a session exists but nothing is in flight — \
+            `continue` to give it another turn), and no such session at all."
     )]
     fn status(&self, Parameters(args): Parameters<StatusArgs>) -> String {
         match session::status(&self.state, &args.name) {

hive-subagent-mcp/src/session.rs

@@ -2,7 +2,8 @@
 //! only while it's alive, and push exactly one todo when it finishes.
 //!
 //! **No task files, no restart recovery.** The daemon's only state is an
-//! in-memory `name -> Cancel` map, live for exactly as long as the process
+//! in-memory `name -> Option<Cancel>` map (see `State`'s own doc for what
+//! the `None`/`Some` split is for), live for exactly as long as the process
 //! is — a daemon restart means whatever was running gets killed with it
 //! (`tokio`'s own child-process drop semantics), not adopted. The durable
 //! record of a subagent's existence is `hive_claude::SessionStore` — claude's
@@ -28,12 +29,20 @@ use std::sync::{Arc, Mutex, PoisonError};
 
 use hive_claude::{Attach, Cancel, Claude, Config, NoopSink, SessionStore};
 
-/// This daemon's whole state: which names currently have a live process,
-/// and where to push the completion todo. `Arc`-wrapped so the background
-/// task that drives a turn to completion can outlive the tool call that
-/// started it.
+/// This daemon's whole state: which names currently have a live process or
+/// a reservation in flight, and where to push the completion todo.
+/// `Arc`-wrapped so the background task that drives a turn to completion
+/// can outlive the tool call that started it.
+///
+/// The map value is `Option<Cancel>`: `None` means `name` is reserved for
+/// an in-flight `start`/`continue` that hasn't reached a confirmed
+/// `Claude::spawn` yet; `Some(cancel)` means a real process is tracked and
+/// interruptible. The `None` state exists to close a real TOCTOU window a
+/// reviewer caught in the original check-then-insert version: checking "is
+/// `name` free" and committing to it are two different lock acquisitions
+/// unless the check *is* the reservation — see `reserve`.
 pub struct State {
-    running: Mutex<HashMap<String, Cancel>>,
+    running: Mutex<HashMap<String, Option<Cancel>>>,
     socket: PathBuf,
 }
 
@@ -46,11 +55,42 @@ impl State {
         }
     }
 
-    fn is_running(&self, name: &str) -> bool {
+    /// `Some(true)` — a live process is tracked, interruptible. `Some(false)`
+    /// — reserved for an in-flight start/continue, not yet a confirmed
+    /// spawn. `None` — nothing tracked under `name` at all.
+    fn occupancy(&self, name: &str) -> Option<bool> {
         self.running
             .lock()
             .unwrap_or_else(PoisonError::into_inner)
-            .contains_key(name)
+            .get(name)
+            .map(Option::is_some)
+    }
+
+    /// Atomically claim `name` for an in-flight start/continue: check
+    /// "is anything tracked under this name" and "commit to this call
+    /// owning it" in the *same* lock acquisition, so two calls racing the
+    /// same name can't both pass a check before either commits (the exact
+    /// same-name concurrent-run hazard this module's doc warns about).
+    /// Returns `false` (reserving nothing) if `name` is already reserved
+    /// or running.
+    fn reserve(&self, name: &str) -> bool {
+        let mut running = self.running.lock().unwrap_or_else(PoisonError::into_inner);
+        if running.contains_key(name) {
+            return false;
+        }
+        running.insert(name.to_owned(), None);
+        true
+    }
+
+    /// Release a reservation that never made it to a real spawn (an error
+    /// on the slow path between `reserve` and `Claude::spawn` succeeding).
+    /// A no-op if the entry was already upgraded to `Some` — this only ever
+    /// clears a still-`None` placeholder, never a live process.
+    fn release_reservation(&self, name: &str) {
+        let mut running = self.running.lock().unwrap_or_else(PoisonError::into_inner);
+        if matches!(running.get(name), Some(None)) {
+            running.remove(name);
+        }
     }
 }
 
@@ -129,9 +169,27 @@ pub fn start(
     trigger: String,
 ) -> anyhow::Result<String> {
     validate_name(name)?;
-    if state.is_running(name) {
+    if !state.reserve(name) {
         anyhow::bail!("subagent `{name}` is already running — use `continue` or `interrupt`");
     }
+    let result = start_reserved(state, name, model, prompt_file, trigger);
+    if result.is_err() {
+        state.release_reservation(name);
+    }
+    result
+}
+
+/// The slow, fallible part of `start`, run only after `reserve` has
+/// already closed the TOCTOU window — split out so `start` can release the
+/// reservation on any error path here without duplicating that logic per
+/// failure site.
+fn start_reserved(
+    state: &Arc<State>,
+    name: &str,
+    model: Option<String>,
+    prompt_file: &str,
+    trigger: String,
+) -> anyhow::Result<String> {
     let config = build_config(name, model, Some(prompt_file));
     let store = build_store(&config)?;
     if store.find_by_title(name).is_some() {
@@ -169,11 +227,27 @@ pub fn continue_(
     model: Option<String>,
 ) -> anyhow::Result<String> {
     validate_name(name)?;
-    if state.is_running(name) {
+    if !state.reserve(name) {
         anyhow::bail!(
             "subagent `{name}` is already running — use `interrupt` first if you meant to redirect it"
         );
     }
+    let result = continue_reserved(state, name, prompt, model);
+    if result.is_err() {
+        state.release_reservation(name);
+    }
+    result
+}
+
+/// The slow, fallible part of `continue_`, run only after `reserve` has
+/// already closed the TOCTOU window — same split rationale as
+/// `start_reserved`.
+fn continue_reserved(
+    state: &Arc<State>,
+    name: &str,
+    prompt: String,
+    model: Option<String>,
+) -> anyhow::Result<String> {
     let config = build_config(name, model, None);
     let store = build_store(&config)?;
     if store.find_by_title(name).is_none() {
@@ -207,11 +281,14 @@ fn spawn_and_track(
 ) -> anyhow::Result<String> {
     let running = Claude::spawn(config, attach)
         .map_err(|e| anyhow::anyhow!("starting the subagent process failed: {e}"))?;
+    // Upgrades the `None` reservation `reserve` already placed here to a
+    // real cancel handle — same key, so there's no window where `name`
+    // reads as unoccupied between the reservation and this insert.
     state
         .running
         .lock()
         .unwrap_or_else(PoisonError::into_inner)
-        .insert(name.to_owned(), running.cancel_handle());
+        .insert(name.to_owned(), Some(running.cancel_handle()));
 
     let state = Arc::clone(state);
     let task_name = name.to_owned();
@@ -236,17 +313,24 @@ fn spawn_and_track(
 }
 
 /// Report whether `name` is currently running — a zero-cost check that
-/// never launches a process, unlike `continue`. Distinguishes three
-/// states: running, idle (a session exists but nothing is in flight), and
-/// no such session at all.
+/// never launches a process, unlike `continue`. Distinguishes four states:
+/// running, starting (reserved, not yet a confirmed spawn — see `State`'s
+/// doc), idle (a session exists but nothing is in flight), and no such
+/// session at all.
 ///
 /// # Errors
 ///
 /// An invalid name, or no session — running or on disk — under `name`.
 pub fn status(state: &State, name: &str) -> anyhow::Result<String> {
     validate_name(name)?;
-    if state.is_running(name) {
-        return Ok(format!("subagent `{name}` is running"));
+    match state.occupancy(name) {
+        Some(true) => return Ok(format!("subagent `{name}` is running")),
+        Some(false) => {
+            return Ok(format!(
+                "subagent `{name}` is starting — not yet confirmed running"
+            ));
+        }
+        None => {}
     }
     let config = build_config(name, None, None);
     let store = build_store(&config)?;
@@ -260,21 +344,33 @@ pub fn status(state: &State, name: &str) -> anyhow::Result<String> {
 }
 
 /// Signal `name`'s running process — `force` picks SIGKILL over SIGINT (see
-/// `hive_claude::Cancel::cancel`). Refuses a name with nothing running:
-/// there's no queued/pending state to cancel pre-emptively any more (see
-/// the module doc), only "running" or "not tracked."
+/// `hive_claude::Cancel::cancel`). Refuses a name with nothing running: no
+/// entry at all, or one still in the brief `reserve`d-but-not-yet-spawned
+/// window (nothing to signal yet — the reservation is put back so a
+/// concurrent `start`/`continue` for the same name still gets refused).
 ///
 /// # Errors
 ///
-/// An invalid name, or nothing currently running under `name`.
+/// An invalid name, nothing tracked under `name`, or `name` is still
+/// starting (reserved, not yet a confirmed spawn).
 pub fn interrupt(state: &State, name: &str, force: bool) -> anyhow::Result<String> {
     validate_name(name)?;
     let mut running = state.running.lock().unwrap_or_else(PoisonError::into_inner);
-    let Some(cancel) = running.remove(name) else {
-        anyhow::bail!("no subagent named `{name}` is currently running");
-    };
-    cancel.cancel(force);
-    Ok(format!("interrupt sent to subagent `{name}`"))
+    match running.remove(name) {
+        None => anyhow::bail!("no subagent named `{name}` is currently running"),
+        Some(None) => {
+            running.insert(name.to_owned(), None);
+            anyhow::bail!(
+                "subagent `{name}` is still starting — not yet confirmed running, try again \
+                 shortly"
+            );
+        }
+        Some(Some(cancel)) => {
+            drop(running);
+            cancel.cancel(force);
+            Ok(format!("interrupt sent to subagent `{name}`"))
+        }
+    }
 }
 
 /// Push `name`'s one-shot completion todo. Best-effort: a connect/write
@@ -298,6 +394,52 @@ async fn push_completion_todo(socket: &std::path::Path, name: &str, summary: &st
 mod tests {
     use super::*;
 
+    // `Cancel` can only be constructed from a real spawned process (no test
+    // fixture in `hive_claude` for it), so these exercise the reservation
+    // half of `State` directly — the actual TOCTOU-closure logic — rather
+    // than the full start/spawn path.
+
+    #[test]
+    fn reserve_is_exclusive_for_the_same_name() {
+        let state = State::new(PathBuf::from("/dev/null"));
+        assert!(state.reserve("dup"), "first reservation should succeed");
+        assert!(
+            !state.reserve("dup"),
+            "a second reservation for the same name must be refused — this is the exact race \
+             argus found: two calls both passing a check before either commits"
+        );
+    }
+
+    #[test]
+    fn reserve_does_not_cross_block_different_names() {
+        let state = State::new(PathBuf::from("/dev/null"));
+        assert!(state.reserve("a"));
+        assert!(state.reserve("b"), "unrelated names must not block each other");
+    }
+
+    #[test]
+    fn release_reservation_frees_the_name_for_reuse() {
+        let state = State::new(PathBuf::from("/dev/null"));
+        assert!(state.reserve("n"));
+        state.release_reservation("n");
+        assert!(
+            state.reserve("n"),
+            "releasing a still-`None` reservation must free the name again"
+        );
+    }
+
+    #[test]
+    fn occupancy_reflects_the_reserved_but_not_running_state() {
+        let state = State::new(PathBuf::from("/dev/null"));
+        assert_eq!(state.occupancy("never-reserved"), None);
+        state.reserve("n");
+        assert_eq!(
+            state.occupancy("n"),
+            Some(false),
+            "reserved-but-not-yet-spawned must read as occupied-but-not-running"
+        );
+    }
+
     #[test]
     fn otel_attrs_appends_when_ambient_var_is_set() {
         // SAFETY: test-only env mutation, single-threaded within this fn's

nix/agent-modules/mcp.nix

@@ -339,9 +339,9 @@ in
       serviceConfig = {
         ExecStart = "${config.hyperhive.packages.hive-subagent-daemon}/bin/hive-subagent-daemon --http 127.0.0.1:${toString config.hyperhive.mcp.subagentHttpPort}";
         SyslogIdentifier = "hive-subagent-daemon";
-        # `always`, same reasoning as `hive-bash-daemon`: the MCP tool is
+        # `always`, same reasoning as `hive-bash-daemon`: the MCP tools are
         # served in-process, so a down window is total loss of
-        # `spawn_subagent` with no stdio fallback.
+        # `start`/`continue`/`status`/`interrupt` with no stdio fallback.
         Restart = "always";
         RestartSec = 3;
         User = userName;