hyperhive/swarm-controller/src/agent_login.rs

//! Whether each agent last said its claude needs a login, for
//! `GET /api/agents/status`'s `needs_login`.
//!
//! The flag comes from the agent's own turn-state header
//! (`$SWARM.agent-state.<agent>`), not from its hive. Agents publish a header
//! when it changes and re-send it once a minute, and the queue keeps none, so
//! an agent has no entry for up to a minute after this daemon starts — the
//! same gap [`crate::inference::Reported`] has on the same subject.

use std::collections::HashMap;
use std::sync::{Arc, Mutex, PoisonError};
use std::time::Duration;

use futures_util::StreamExt as _;
use swarm_wire::agent_login::Login;

/// The latest login flag each agent reported, by agent.
#[derive(Debug, Default)]
pub(crate) struct Reported {
    latest: Mutex<HashMap<String, bool>>,
}

impl Reported {
    /// Start recording every agent's login flag from `client`.
    pub(crate) fn spawn(client: async_nats::Client) -> Arc<Self> {
        let reported = Arc::new(Self::default());
        let recording = Arc::clone(&reported);
        tokio::spawn(async move { recording.follow(&client).await });
        reported
    }

    /// Subscribe to every agent's header, and subscribe again if the
    /// subscription ends.
    async fn follow(&self, client: &async_nats::Client) {
        const RETRY: Duration = Duration::from_secs(5);
        let subject = swarm_queue_client::agent_state_subject_wildcard();
        loop {
            match client.subscribe(subject.clone()).await {
                Ok(mut headers) => {
                    tracing::info!(%subject, "agent login: following agent turn-state headers");
                    while let Some(msg) = headers.next().await {
                        self.record(&msg.subject, &msg.payload);
                    }
                    tracing::warn!(%subject, "agent login: header subscription ended");
                }
                Err(e) => tracing::warn!(%subject, error = %e, "agent login: subscribe failed"),
            }
            tokio::time::sleep(RETRY).await;
        }
    }

    /// Keep one header's login flag. A header without it, from an agent that
    /// does not report it, removes what was kept.
    fn record(&self, subject: &str, payload: &[u8]) {
        let Some(agent) = swarm_queue_client::agent_from_state_subject(subject) else {
            return;
        };
        let mut latest = self.latest.lock().unwrap_or_else(PoisonError::into_inner);
        match serde_json::from_slice::<Login>(payload) {
            Ok(Login {
                needs_login: Some(needs_login),
            }) => {
                latest.insert(agent.to_owned(), needs_login);
            }
            Ok(Login { needs_login: None }) => {
                latest.remove(agent);
            }
            Err(e) => {
                tracing::debug!(%agent, error = %e, "agent login: header did not decode");
                latest.remove(agent);
            }
        }
    }

    /// `agent`'s latest flag; `None` when no header since this daemon
    /// started has carried one.
    pub(crate) fn get(&self, agent: &str) -> Option<bool> {
        self.latest
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .get(agent)
            .copied()
    }
}

#[cfg(test)]
mod tests {
    use super::Reported;

    fn header(needs_login: Option<bool>) -> Vec<u8> {
        let mut header = serde_json::json!({
            "turn_state": "idle",
            "turn_state_since": "2026-09-13T12:35:03Z",
            "agent_state": "up",
            "model": "sonnet",
            "effort": "high",
            "model_source": "default",
            "effort_source": "default",
            "resolved_model": null,
            "context_window_tokens": 200_000,
            "ctx_usage": null,
            "cost_usage": null,
        });
        if let Some(needs_login) = needs_login {
            header["needs_login"] = needs_login.into();
        }
        serde_json::to_vec(&header).unwrap()
    }

    #[test]
    fn a_header_without_the_flag_reports_nothing() {
        let reported = Reported::default();
        let subject = swarm_queue_client::agent_state_subject("iris");
        reported.record(&subject, &header(None));
        assert_eq!(reported.get("iris"), None);
    }

    #[test]
    fn the_latest_header_wins_and_one_without_the_flag_clears_it() {
        let reported = Reported::default();
        let subject = swarm_queue_client::agent_state_subject("iris");
        reported.record(&subject, &header(Some(true)));
        assert_eq!(reported.get("iris"), Some(true));
        reported.record(&subject, &header(Some(false)));
        assert_eq!(reported.get("iris"), Some(false));
        reported.record(&subject, &header(Some(true)));
        reported.record(&subject, &header(None));
        assert_eq!(reported.get("iris"), None);
        reported.record(&subject, &header(Some(true)));
        reported.record(&subject, b"not json");
        assert_eq!(reported.get("iris"), None);
    }

    #[test]
    fn a_header_is_kept_under_the_agent_its_subject_names() {
        let reported = Reported::default();
        reported.record(
            &swarm_queue_client::agent_state_subject("iris"),
            &header(Some(true)),
        );
        reported.record(
            &swarm_queue_client::term_subject("atlas"),
            &header(Some(true)),
        );
        assert_eq!(reported.get("iris"), Some(true));
        assert_eq!(reported.get("atlas"), None);
    }
}