hive-subagent-mcp: publish each subagent's run state to the swarm queue

hive-subagent-mcp/src/session.rs

@@ -738,6 +738,9 @@ impl State {
     /// here is what makes a signal for a session that already ended a 404
     /// rather than a late stop reason recorded against it.
     fn finish_turn(&self, name: &str, end: &TurnEnd) {
+        if let Some(term) = &self.term {
+            term.finished(name, exit_status(end));
+        }
         self.revoke_signal_token(name);
         let mut killed = self.killed.lock().unwrap_or_else(PoisonError::into_inner);
         match end {
@@ -755,6 +758,15 @@ impl State {
             .remove(name);
     }
 
+    /// Tell the swarm `name`'s run began, with the model and effort `info`
+    /// names. The end is reported by `finish_turn`, which every run's end goes
+    /// through; the gaps between a goal run's turns are not an end.
+    fn announce_running(&self, name: &str, info: &Arc<crate::swarm_term::RunInfo>) {
+        if let Some(term) = &self.term {
+            term.running(name, info);
+        }
+    }
+
     /// The signal `name`'s last turn died on, if it died on one.
     fn killed_by(&self, name: &str) -> Option<i32> {
         self.killed
@@ -1853,6 +1865,19 @@ impl LivenessSink {
     }
 }
 
+/// The exit status a run's end reports to the swarm: `0` when it completed,
+/// nothing when a signal killed it, and `1` for any other failure. A failed
+/// turn's own exit code is not kept on [`TurnEnd`] (it carries the message
+/// shown to the caller), and an ACP failure has no process status at all, so
+/// the swarm sees "failed" and the end todo has the detail.
+fn exit_status(end: &TurnEnd) -> Option<i32> {
+    match end {
+        TurnEnd::Complete => Some(0),
+        TurnEnd::Killed { .. } => None,
+        TurnEnd::Failed(_) => Some(1),
+    }
+}
+
 /// What a run's rows say about how it was started: its `config`'s model and
 /// effort, which the swarm shows beside the subagent's name.
 fn run_info(config: &Config) -> Arc<crate::swarm_term::RunInfo> {
@@ -1924,6 +1949,7 @@ fn spawn_and_track(
     };
 
     let info = run_info(config);
+    state.announce_running(name, &info);
     let config = config.clone();
     let state = Arc::clone(state);
     let task_name = name.to_owned();
@@ -2040,6 +2066,7 @@ fn spawn_and_track_acp(
     state.note_event(name);
 
     let info = run_info(&config);
+    state.announce_running(name, &info);
     let state = Arc::clone(state);
     let task_name = name.to_owned();
     tokio::spawn(async move {
@@ -3164,6 +3191,15 @@ mod tests {
         }
     }
 
+    /// What the swarm is told when a run ends: success is 0, a signal is no
+    /// code at all, and any other failure is 1.
+    #[test]
+    fn a_runs_end_reports_its_exit_status_to_the_swarm() {
+        assert_eq!(exit_status(&TurnEnd::Complete), Some(0));
+        assert_eq!(exit_status(&TurnEnd::Killed { signal: 9 }), None);
+        assert_eq!(exit_status(&TurnEnd::Failed("boom".to_owned())), Some(1));
+    }
+
     #[test]
     fn classify_end_separates_a_signalled_child_from_a_clean_one() {
         assert_eq!(classify_end(Ok(()), None), TurnEnd::Complete);

hive-subagent-mcp/src/swarm_term.rs

@@ -22,6 +22,9 @@
 //! row published while that stream does not exist yet still reaches a live
 //! subscriber, but the swarm does not list the subagent from it.
 //!
+//! **Run state too.** The same task publishes each subagent's running or
+//! finished state to `$KV.subagent-status.<agent>.<subagent>`, with a heartbeat.
+//!
 //! **Best-effort.** A subagent's run never waits on the queue: no store, a
 //! refused credential, a full queue and a failed publish are each a log
 //! line, and the run goes on.
@@ -38,6 +41,7 @@ use swarm_secret_client::{
     client::{DEFAULT_CERT_MOUNT, ENV_ADDR, ENV_CACERT, Settings},
     policy, queue,
 };
+use swarm_wire::subagent_status::{self, State as RunState, Status};
 use swarm_wire::subagent_term;
 use tokio::sync::mpsc;
 
@@ -114,6 +118,20 @@ fn header_value(value: Option<&str>) -> Option<async_nats::HeaderValue> {
     value.parse().ok()
 }
 
+/// A change in whether a subagent is running, stamped when it happened.
+enum StatusEvent {
+    Running {
+        subagent: String,
+        info: Arc<RunInfo>,
+        at: i64,
+    },
+    Finished {
+        subagent: String,
+        exit: Option<i32>,
+        at: i64,
+    },
+}
+
 struct Line {
     subagent: String,
     info: Arc<RunInfo>,
@@ -125,6 +143,9 @@ struct Line {
 /// The sink's handle on the publish task.
 pub struct Publisher {
     tx: mpsc::Sender<Line>,
+    /// Unbounded, unlike `tx`: a run sends two of these, and a lost
+    /// `Finished` would leave the task heartbeating a subagent that ended.
+    status: mpsc::UnboundedSender<StatusEvent>,
     dropped: Arc<AtomicU64>,
 }
 
@@ -158,27 +179,53 @@ impl Publisher {
             stream = %subagent_term::stream_name(&target.agent),
             "publishing subagent terminals to the swarm queue as this agent"
         );
-        let (publisher, rx) = Self::channel(QUEUE_DEPTH);
+        let (publisher, rx, status_rx) = Self::channel(QUEUE_DEPTH);
-        tokio::spawn(run(rx, target, Arc::clone(&publisher.dropped)));
+        tokio::spawn(run(rx, status_rx, target, Arc::clone(&publisher.dropped)));
         Some(publisher)
     }
 
-    fn channel(depth: usize) -> (Self, mpsc::Receiver<Line>) {
+    fn channel(
+        depth: usize,
+    ) -> (
+        Self,
+        mpsc::Receiver<Line>,
+        mpsc::UnboundedReceiver<StatusEvent>,
+    ) {
         let (tx, rx) = mpsc::channel(depth);
+        let (status, status_rx) = mpsc::unbounded_channel();
         let publisher = Self {
             tx,
+            status,
             dropped: Arc::new(AtomicU64::new(0)),
         };
-        (publisher, rx)
+        (publisher, rx, status_rx)
+    }
+
+    /// `subagent` started a run, on the model and effort `info` names.
+    pub fn running(&self, subagent: &str, info: &Arc<RunInfo>) {
+        // A closed receiver means the task is gone; there is nobody to tell.
+        let _ = self.status.send(StatusEvent::Running {
+            subagent: subagent.to_owned(),
+            info: Arc::clone(info),
+            at: unix_now(),
+        });
+    }
+
+    /// `subagent`'s run ended, with `exit` as its process's exit status, or
+    /// `None` when a signal killed it.
+    pub fn finished(&self, subagent: &str, exit: Option<i32>) {
+        let _ = self.status.send(StatusEvent::Finished {
+            subagent: subagent.to_owned(),
+            exit,
+            at: unix_now(),
+        });
     }
 
     /// Queue `output` from `subagent` for publishing, with the `info` its rows
     /// carry as headers. Never blocks: a full queue drops the line and counts
     /// it.
     pub fn offer(&self, subagent: &str, info: &Arc<RunInfo>, output: Output) {
-        let ts = SystemTime::now()
+        let ts = unix_now();
-            .duration_since(UNIX_EPOCH)
-            .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(i64::MAX));
         let line = Line {
             subagent: subagent.to_owned(),
             info: Arc::clone(info),
@@ -191,6 +238,13 @@ impl Publisher {
     }
 }
 
+/// Unix seconds now.
+fn unix_now() -> i64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(i64::MAX))
+}
+
 /// Where to publish, and what to connect with.
 pub(crate) struct Target {
     pub(crate) url: String,
@@ -321,31 +375,130 @@ impl Queue {
     }
 }
 
-async fn run(mut rx: mpsc::Receiver<Line>, target: Target, dropped: Arc<AtomicU64>) {
+async fn run(
+    mut rx: mpsc::Receiver<Line>,
+    mut status_rx: mpsc::UnboundedReceiver<StatusEvent>,
+    target: Target,
+    dropped: Arc<AtomicU64>,
+) {
     let target = Arc::new(target);
     let mut ctxs: HashMap<String, ClassifyCtx> = HashMap::new();
     let mut queue = Queue::default();
-    while let Some(line) = rx.recv().await {
+    let mut runs: HashMap<String, Status> = HashMap::new();
-        let missed = dropped.swap(0, Ordering::Relaxed);
+    let period = Duration::from_secs(subagent_status::HEARTBEAT_S.unsigned_abs());
-        if missed > 0 {
+    let mut beat = tokio::time::interval_at(tokio::time::Instant::now() + period, period);
-            tracing::warn!(missed, "subagent terminal: queue full, lines dropped");
+    loop {
+        tokio::select! {
+            line = rx.recv() => {
+                let Some(line) = line else { return };
+                let missed = dropped.swap(0, Ordering::Relaxed);
+                if missed > 0 {
+                    tracing::warn!(missed, "subagent terminal: queue full, lines dropped");
+                }
+                publish_line(&line, &mut ctxs, &mut queue, &target).await;
+            }
+            Some(event) = status_rx.recv() => {
+                let (subagent, status) = apply(&mut runs, event);
+                publish_status(&mut queue, &target, &subagent, &status).await;
+            }
+            _ = beat.tick() => {
+                for (subagent, status) in refresh(&mut runs, unix_now()) {
+                    publish_status(&mut queue, &target, &subagent, &status).await;
+                }
+            }
         }
-        let ctx = ctxs.entry(line.subagent.clone()).or_default();
+    }
-        let rows = classify(&line.output, line.ts, ctx);
+}
-        if rows.is_empty() {
+
-            continue;
+/// Publish the rows one line of output makes.
+async fn publish_line(
+    line: &Line,
+    ctxs: &mut HashMap<String, ClassifyCtx>,
+    queue: &mut Queue,
+    target: &Arc<Target>,
+) {
+    let ctx = ctxs.entry(line.subagent.clone()).or_default();
+    let rows = classify(&line.output, line.ts, ctx);
+    if rows.is_empty() {
+        return;
+    }
+    let Some(client) = queue.ready(target).await else {
+        return;
+    };
+    let subject = target.subject(&line.subagent);
+    let headers = line.info.headers();
+    for row in rows {
+        publish(&client, &subject, &headers, row).await;
+    }
+}
+
+/// Fold `event` into the runs the task is keeping alive, and return the status
+/// to publish for it. A finished subagent leaves the map, so it is not
+/// heartbeated again; its model and effort carry over to the final status.
+fn apply(runs: &mut HashMap<String, Status>, event: StatusEvent) -> (String, Status) {
+    match event {
+        StatusEvent::Running { subagent, info, at } => {
+            let status = Status {
+                state: RunState::Running,
+                since: at,
+                exit: None,
+                model: info.model.clone(),
+                effort: info.effort.clone(),
+                ts: at,
+            };
+            runs.insert(subagent.clone(), status.clone());
+            (subagent, status)
         }
-        let Some(client) = queue.ready(&target).await else {
+        StatusEvent::Finished { subagent, exit, at } => {
-            continue;
+            let previous = runs.remove(&subagent);
-        };
+            let status = Status {
-        let subject = target.subject(&line.subagent);
+                state: RunState::Finished,
-        let headers = line.info.headers();
+                since: at,
-        for row in rows {
+                exit,
-            publish(&client, &subject, &headers, row).await;
+                model: previous.as_ref().and_then(|p| p.model.clone()),
+                effort: previous.and_then(|p| p.effort),
+                ts: at,
+            };
+            (subagent, status)
         }
     }
 }
 
+/// Stamp every running subagent's status with `ts`, the heartbeat, and return
+/// them to publish.
+fn refresh(runs: &mut HashMap<String, Status>, ts: i64) -> Vec<(String, Status)> {
+    runs.iter_mut()
+        .map(|(subagent, status)| {
+            status.ts = ts;
+            (subagent.clone(), status.clone())
+        })
+        .collect()
+}
+
+/// Publish `status` as `subagent`'s value in the status bucket. A KV put is a
+/// publish to the key's subject, which the agent's credential may do and the
+/// controller's bucket captures. Every failure is a log line.
+async fn publish_status(queue: &mut Queue, target: &Arc<Target>, subagent: &str, status: &Status) {
+    let Some(client) = queue.ready(target).await else {
+        return;
+    };
+    if let Err(e) = swarm_queue_client::ensure_connected(&client) {
+        tracing::warn!(error = %swarm_queue_client::chain(&e), "subagent status: publish skipped");
+        return;
+    }
+    let payload = match serde_json::to_vec(status) {
+        Ok(payload) => payload,
+        Err(e) => {
+            tracing::warn!(error = %e, "subagent status: serialising failed, not published");
+            return;
+        }
+    };
+    let subject = subagent_status::subject(&target.agent, subagent);
+    if let Err(e) = client.publish(subject, payload.into()).await {
+        tracing::warn!(error = %e, subagent, "subagent status: publish failed");
+    }
+}
+
 /// The rows one line of output renders as, each stamped with `ts`.
 fn classify(output: &Output, ts: i64, ctx: &mut ClassifyCtx) -> Vec<TermMsg> {
     let rows = match output {
@@ -463,13 +616,80 @@ mod tests {
     /// and counts it.
     #[test]
     fn a_full_queue_drops_and_counts_rather_than_blocks() {
-        let (publisher, _rx) = Publisher::channel(1);
+        let (publisher, _rx, _status) = Publisher::channel(1);
         let info = Arc::new(RunInfo::default());
         publisher.offer("scout", &info, Output::Stdout("one".to_owned()));
         publisher.offer("scout", &info, Output::Stdout("two".to_owned()));
         assert_eq!(publisher.dropped.load(Ordering::Relaxed), 1);
     }
 
+    fn event_running(subagent: &str, at: i64) -> StatusEvent {
+        StatusEvent::Running {
+            subagent: subagent.to_owned(),
+            info: Arc::new(RunInfo {
+                model: Some("opus".to_owned()),
+                effort: None,
+            }),
+            at,
+        }
+    }
+
+    /// A start is `running` since then with the run's model; the end is
+    /// `finished` with the exit status and the same model, and the subagent
+    /// is not heartbeated afterwards.
+    #[test]
+    fn a_run_is_running_then_finished_with_its_exit_and_model() {
+        let mut runs = HashMap::new();
+        let (name, started) = apply(&mut runs, event_running("scout", 100));
+        assert_eq!(name, "scout");
+        assert_eq!(started.state, RunState::Running);
+        assert_eq!((started.since, started.ts, started.exit), (100, 100, None));
+        assert_eq!(started.model.as_deref(), Some("opus"));
+
+        let finished = StatusEvent::Finished {
+            subagent: "scout".to_owned(),
+            exit: Some(0),
+            at: 160,
+        };
+        let (_, ended) = apply(&mut runs, finished);
+        assert_eq!(ended.state, RunState::Finished);
+        assert_eq!((ended.since, ended.ts, ended.exit), (160, 160, Some(0)));
+        assert_eq!(ended.model.as_deref(), Some("opus"));
+        assert!(refresh(&mut runs, 220).is_empty(), "nothing left to beat");
+    }
+
+    /// The heartbeat moves `ts` and nothing else, so a reader tells a live
+    /// writer from a dead one by `ts` while `since` still says when it began.
+    #[test]
+    fn the_heartbeat_refreshes_ts_and_keeps_since() {
+        let mut runs = HashMap::new();
+        apply(&mut runs, event_running("scout", 100));
+        apply(&mut runs, event_running("other", 130));
+        let mut beat = refresh(&mut runs, 160);
+        beat.sort_by(|a, b| a.0.cmp(&b.0));
+        assert_eq!(beat.len(), 2);
+        assert_eq!((beat[0].0.as_str(), beat[0].1.since, beat[0].1.ts), ("other", 130, 160));
+        assert_eq!((beat[1].0.as_str(), beat[1].1.since, beat[1].1.ts), ("scout", 100, 160));
+        assert!(beat.iter().all(|(_, s)| s.state == RunState::Running));
+    }
+
+    /// A lost end would leave the task heartbeating a subagent that ended, so
+    /// the status events are never dropped for a full queue.
+    #[test]
+    fn status_events_are_not_bounded_by_the_line_queue() {
+        let (publisher, _rx, mut status) = Publisher::channel(1);
+        let info = Arc::new(RunInfo::default());
+        for _ in 0..10 {
+            publisher.running("scout", &info);
+            publisher.finished("scout", Some(0));
+        }
+        let mut seen = 0;
+        while status.try_recv().is_ok() {
+            seen += 1;
+        }
+        assert_eq!(seen, 20);
+    }
+
     /// The list endpoint reads these two headers off a subagent's newest row.
     #[test]
     fn a_row_carries_the_model_and_effort_the_subagent_runs() {