hyperhive/swarm-controller/src/agent_status.rs

//! Swarm-wide view of what each agent last said about itself.
//!
//! Sibling of [`crate::status`], not a mode of it — see
//! `swarm_wire::agent_status`'s module doc for why the two buckets
//! stay separate. Everything about *how* a snapshot is read (queue is the
//! store, no cache, freshness derived at read time, rows come from the
//! roster so an empty bucket cannot render as "every agent is fine") is
//! identical to [`crate::status`]; this module differs only in whose
//! roster it reads and how it recovers a hive name for a reporting agent.
//!
//! **The roster here is agent *names* only** ([`crate::auth::AuthBridge::list_agent_identities`]) —
//! it does not say which hive an agent lives on. That is not a gap this
//! module needs to close: a *reporting* agent's hive comes straight out of
//! its bucket key (`swarm_wire::agent_status::split_key`), and a
//! *never-reported* agent has no bucket key, so [`AgentStatusRow::hive`]
//! is `None` out of this module in exactly that one case. The
//! `GET /api/agents/status` handler fills it from the wanted declaration
//! that places the agent.

use std::collections::BTreeMap;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use anyhow::{Context, Result};
use futures_util::TryStreamExt as _;
use serde::{Deserialize, Serialize};
use utoipa::ToSchema;

use crate::status::Freshness;

/// One row of the aggregate.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, ToSchema)]
pub struct AgentStatusRow {
    pub name: String,
    /// The hive this agent reported from, taken from its bucket key.
    /// `None` when the agent has never reported. Out of this reader only:
    /// `GET /api/agents/status` fills it from the declaration that places
    /// the agent.
    pub hive: Option<String>,
    pub freshness: Freshness,
    /// When the snapshot was stored, unix seconds — the bucket's own
    /// stamp, same rule as [`crate::status::HiveStatus::last_seen_unix`].
    pub last_seen_unix: Option<i64>,
    pub age_seconds: Option<u64>,
    /// The agent's status text + running flag, unopened. `None` exactly
    /// when `freshness` is `NeverReported`, or when the stored bytes
    /// failed to decode as [`swarm_wire::agent_status::AgentStatus`]
    /// (logged as a warning naming the agent).
    pub snapshot: Option<serde_json::Value>,
    /// The agent's open config PR, if any. Always `None` coming out of
    /// [`AgentStatusReader::view`] itself — this module has no forge
    /// client and stays that way, same separation
    /// [`crate::status`]/[`crate::agent_status`] already keep from each
    /// other. `GET /api/agents/status`'s handler fills this field in
    /// after the fact from `AppState::config_prs`, which is the one
    /// place that already holds both a roster-shaped answer and a
    /// config-PR cache — see the handler for why merging there, not
    /// here, is what makes this the single call swarm-ui's agent roster
    /// page needs.
    pub config_pr: Option<crate::forge::ConfigPrStatus>,
    /// The agent's declared wanted state (`"up"`/`"offline"`), if this
    /// hive has one on record. Same rule as `config_pr`: always `None`
    /// out of this reader, filled in by the `GET /api/agents/status`
    /// handler from `AppState::wanted` — this module has no notion of a
    /// *declaration* (a swarm-level intent), only of what an agent last
    /// *reported about itself*, and merging a second bucket's read in
    /// here would blur that boundary for the same reason `config_pr`
    /// doesn't merge forge state in directly.
    pub wanted: Option<String>,
    /// Whether the agent's latest turn-state header said its claude needs a
    /// login. `None` when no header since swarm-controller started carried
    /// the flag. Same rule as `config_pr`: always `None` out of this reader,
    /// filled in by the `GET /api/agents/status` handler from
    /// `AppState::agent_login` — this reader sees the hive's bucket, and the
    /// flag comes from the agent itself.
    pub needs_login: Option<bool>,
}

/// A snapshot as retained by the bucket, with the hive it was published
/// under already split out of the key.
#[derive(Clone, Debug)]
struct Offered {
    hive: String,
    received_at: SystemTime,
    payload: Option<serde_json::Value>,
}

/// Reads the per-agent status aggregate out of the KV bucket. Same shape
/// as [`crate::status::StatusReader`] — see that type for why the bucket
/// is resolved lazily and cached, and why a resolution failure is not.
pub struct AgentStatusReader {
    client: async_nats::Client,
    store: tokio::sync::OnceCell<async_nats::jetstream::kv::Store>,
    stale_after: Duration,
}

impl AgentStatusReader {
    #[must_use]
    pub fn new(client: async_nats::Client, stale_after: Duration) -> Self {
        Self {
            client,
            store: tokio::sync::OnceCell::new(),
            stale_after,
        }
    }

    async fn store(
        &self,
    ) -> std::result::Result<&async_nats::jetstream::kv::Store, swarm_queue_client::Error> {
        self.store
            .get_or_try_init(|| swarm_queue_client::agent_status::open_or_create(&self.client))
            .await
    }

    /// The aggregate, rendered against `now`.
    ///
    /// Every roster agent produces a row whether or not it has ever
    /// reported; a reporting agent outside the roster produces one too —
    /// same two-sided shape as [`crate::status::StatusReader::view`].
    pub async fn view(&self, roster: &[String], now: SystemTime) -> Result<Vec<AgentStatusRow>> {
        swarm_queue_client::ensure_connected(&self.client)?;

        let store = self.store().await?;

        let mut keys = store
            .keys()
            .await
            .context("listing agent-status bucket keys")?;
        let mut entries: BTreeMap<String, Offered> = BTreeMap::new();
        while let Some(key) = keys
            .try_next()
            .await
            .context("reading the agent-status bucket's key list")?
        {
            let Some((hive, agent)) = swarm_wire::agent_status::split_key(&key) else {
                tracing::warn!(%key, "agent-status bucket key doesn't split into hive/agent, skipping");
                continue;
            };
            let Some(entry) = store
                .entry(&key)
                .await
                .with_context(|| format!("reading agent-status entry {key}"))?
            else {
                continue;
            };
            let payload = match serde_json::from_slice::<swarm_wire::agent_status::AgentStatus>(
                &entry.value,
            ) {
                Ok(status) => serde_json::to_value(status).ok(),
                Err(e) => {
                    tracing::warn!(
                        %agent, %hive, error = %e,
                        "agent-status snapshot didn't decode; reporting the timestamp without it"
                    );
                    None
                }
            };
            entries.insert(
                agent.to_owned(),
                Offered {
                    hive: hive.to_owned(),
                    received_at: to_system_time(entry.created.unix_timestamp()),
                    payload,
                },
            );
        }

        Ok(render(roster, &entries, now, self.stale_after))
    }
}

fn to_system_time(unix_seconds: i64) -> SystemTime {
    u64::try_from(unix_seconds).map_or(UNIX_EPOCH, |secs| UNIX_EPOCH + Duration::from_secs(secs))
}

/// Turn a roster plus whatever the bucket held into the rendered rows. Pure
/// for the same testing reason as [`crate::status::render`].
fn render(
    roster: &[String],
    entries: &BTreeMap<String, Offered>,
    now: SystemTime,
    stale_after: Duration,
) -> Vec<AgentStatusRow> {
    let mut rows: Vec<AgentStatusRow> = roster
        .iter()
        .map(|name| match entries.get(name) {
            Some(offered) => row(name.clone(), Freshness::Fresh, offered, now, stale_after),
            None => AgentStatusRow {
                name: name.clone(),
                hive: None,
                freshness: Freshness::NeverReported,
                last_seen_unix: None,
                age_seconds: None,
                snapshot: None,
                config_pr: None,
                wanted: None,
                needs_login: None,
            },
        })
        .collect();

    rows.extend(
        entries
            .iter()
            .filter(|(name, _)| !roster.iter().any(|r| r == *name))
            .map(|(name, offered)| {
                row(name.clone(), Freshness::Unknown, offered, now, stale_after)
            }),
    );
    rows
}

/// `default_freshness` is `Unknown` for an agent outside the roster and
/// `Fresh` for one inside it — `row` then downgrades to `Stale` on age,
/// same two-step [`crate::status::row`] uses.
fn row(
    name: String,
    default_freshness: Freshness,
    offered: &Offered,
    now: SystemTime,
    stale_after: Duration,
) -> AgentStatusRow {
    let age = now.duration_since(offered.received_at).ok();
    let freshness = match (default_freshness, age) {
        (Freshness::Unknown, _) => Freshness::Unknown,
        (_, Some(age)) if age > stale_after => Freshness::Stale,
        _ => Freshness::Fresh,
    };
    AgentStatusRow {
        name,
        hive: Some(offered.hive.clone()),
        freshness,
        last_seen_unix: offered
            .received_at
            .duration_since(UNIX_EPOCH)
            .ok()
            .and_then(|d| i64::try_from(d.as_secs()).ok()),
        age_seconds: age.map(|age| age.as_secs()),
        snapshot: offered.payload.clone(),
        config_pr: None,
        wanted: None,
        needs_login: None,
    }
}

#[cfg(test)]
mod tests {
    use super::{Freshness, Offered, render};
    use std::collections::BTreeMap;
    use std::time::{Duration, SystemTime};

    fn roster() -> Vec<String> {
        vec!["iris".to_owned(), "atlas".to_owned()]
    }

    fn entries(rows: &[(&str, &str, SystemTime)]) -> BTreeMap<String, Offered> {
        rows.iter()
            .map(|(agent, hive, at)| {
                (
                    (*agent).to_owned(),
                    Offered {
                        hive: (*hive).to_owned(),
                        received_at: *at,
                        payload: Some(serde_json::json!({ "running": true })),
                    },
                )
            })
            .collect()
    }

    fn t0() -> SystemTime {
        SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000)
    }

    #[test]
    fn an_empty_bucket_renders_every_roster_agent_as_never_reported() {
        let rows = render(&roster(), &BTreeMap::new(), t0(), Duration::from_mins(2));
        assert_eq!(rows.len(), 2);
        assert!(rows.iter().all(|r| r.freshness == Freshness::NeverReported));
        assert!(rows.iter().all(|r| r.hive.is_none()));
    }

    #[test]
    fn a_reporting_agent_gets_its_hive_from_the_bucket_key() {
        let rows = render(
            &roster(),
            &entries(&[("iris", "pr1ma", t0())]),
            t0(),
            Duration::from_mins(2),
        );
        let iris = rows.iter().find(|r| r.name == "iris").expect("present");
        assert_eq!(iris.hive.as_deref(), Some("pr1ma"));
        assert_eq!(iris.freshness, Freshness::Fresh);
    }

    #[test]
    fn a_stale_report_still_carries_its_hive() {
        let rows = render(
            &roster(),
            &entries(&[("iris", "pr1ma", t0())]),
            t0() + Duration::from_mins(10),
            Duration::from_mins(2),
        );
        let iris = rows.iter().find(|r| r.name == "iris").expect("present");
        assert_eq!(iris.freshness, Freshness::Stale);
        assert_eq!(iris.hive.as_deref(), Some("pr1ma"), "stale is not absent");
    }

    #[test]
    fn an_agent_outside_the_roster_is_surfaced_as_unknown() {
        let rows = render(
            &roster(),
            &entries(&[("ghost", "pr1ma", t0())]),
            t0(),
            Duration::from_mins(2),
        );
        assert_eq!(rows.len(), 3, "two roster agents plus the stranger");
        let ghost = rows.iter().find(|r| r.name == "ghost").expect("present");
        assert_eq!(ghost.freshness, Freshness::Unknown);
        assert_eq!(
            ghost.hive.as_deref(),
            Some("pr1ma"),
            "still known from its key"
        );
    }

    #[test]
    fn one_agent_reporting_does_not_vouch_for_a_silent_one() {
        let rows = render(
            &roster(),
            &entries(&[("iris", "pr1ma", t0())]),
            t0(),
            Duration::from_mins(2),
        );
        let atlas = rows.iter().find(|r| r.name == "atlas").expect("present");
        assert_eq!(atlas.freshness, Freshness::NeverReported);
    }
}